| 72 | |
| 73 | template <typename... Args> |
| 74 | port::StatusOr<std::shared_ptr<TypedKernel<Args...>>> LoadKernelOrGetPtr( |
| 75 | StreamExecutor* executor, absl::string_view kernel_name, |
| 76 | absl::string_view ptx, absl::Span<const uint8> cubin_data) { |
| 77 | using KernelPtrCacheKey = |
| 78 | std::tuple<CUcontext, absl::string_view, absl::string_view>; |
| 79 | |
| 80 | static tensorflow::mutex kernel_ptr_cache_mutex( |
| 81 | tensorflow::LINKER_INITIALIZED); |
| 82 | static auto& kernel_ptr_cache GUARDED_BY(kernel_ptr_cache_mutex) = |
| 83 | *new absl::flat_hash_map<KernelPtrCacheKey, |
| 84 | std::shared_ptr<TypedKernel<Args...>>>(); |
| 85 | CUcontext current_context = CurrentContextOrDie(); |
| 86 | |
| 87 | KernelPtrCacheKey kernel_ptr_cache_key{current_context, kernel_name, ptx}; |
| 88 | |
| 89 | tensorflow::mutex_lock lock(kernel_ptr_cache_mutex); |
| 90 | |
| 91 | auto it = kernel_ptr_cache.find(kernel_ptr_cache_key); |
| 92 | if (it == kernel_ptr_cache.end()) { |
| 93 | TF_ASSIGN_OR_RETURN( |
| 94 | std::shared_ptr<TypedKernel<Args...>> loaded, |
| 95 | executor->CreateTypedKernel<Args...>(kernel_name, ptx, cubin_data)); |
| 96 | it = |
| 97 | kernel_ptr_cache.emplace(kernel_ptr_cache_key, std::move(loaded)).first; |
| 98 | } |
| 99 | |
| 100 | CHECK(it != kernel_ptr_cache.end()); |
| 101 | return it->second; |
| 102 | } |
| 103 | |
| 104 | char* GetTFExtraPTXOptions(); |
| 105 |
nothing calls this directly
no test coverage detected