| 435 | } |
| 436 | |
| 437 | Status EagerContext::RemoveFunction(const string& func) { |
| 438 | bool is_last_ref = false; |
| 439 | { |
| 440 | mutex_lock l(cache_mu_); |
| 441 | auto* registered_function = gtl::FindPtrOrNull(registered_functions_, func); |
| 442 | if (registered_function == nullptr) { |
| 443 | return errors::InvalidArgument("Tried to remove non-existent function '", |
| 444 | func, "'."); |
| 445 | } |
| 446 | is_last_ref = registered_function->RefCountIsOne(); |
| 447 | if (is_last_ref) { |
| 448 | for (auto& key : *registered_function->cached_kernel_keys) { |
| 449 | kernel_cache_.erase(key); |
| 450 | } |
| 451 | registered_functions_.erase(func); |
| 452 | } |
| 453 | registered_function->Unref(); |
| 454 | } |
| 455 | if (is_last_ref) { |
| 456 | // TODO(fishx): Remove remote function as well. |
| 457 | return func_lib_def_.RemoveFunction(func); |
| 458 | } |
| 459 | return Status::OK(); |
| 460 | } |
| 461 | |
| 462 | core::RefCountPtr<KernelAndDevice> EagerContext::GetCachedKernel( |
| 463 | Fprint128 cache_key) { |
nothing calls this directly
no test coverage detected