| 961 | } |
| 962 | |
| 963 | bool GpuExecutor::GetSymbol(const string& symbol_name, |
| 964 | ModuleHandle module_handle, void** mem, |
| 965 | size_t* bytes) { |
| 966 | auto lookup_in_module = [&](CUmodule module) { |
| 967 | CHECK(module != nullptr); |
| 968 | return GpuDriver::GetModuleSymbol(context_, module, symbol_name.c_str(), |
| 969 | reinterpret_cast<CUdeviceptr*>(mem), |
| 970 | bytes); |
| 971 | }; |
| 972 | |
| 973 | { // give limited scope to mutex_lock |
| 974 | absl::MutexLock lock{&in_memory_modules_mu_}; |
| 975 | if (static_cast<bool>(module_handle)) { |
| 976 | auto it = gpu_binary_to_module_.find(module_handle.id()); |
| 977 | CHECK(it != gpu_binary_to_module_.end()); |
| 978 | return lookup_in_module(it->second.first); |
| 979 | } |
| 980 | |
| 981 | for (auto &it : gpu_binary_to_module_) { |
| 982 | if (lookup_in_module(it.second.first)) { |
| 983 | return true; |
| 984 | } |
| 985 | } |
| 986 | } |
| 987 | |
| 988 | LOG(INFO) << "Failed to find symbol in any modules: " << symbol_name; |
| 989 | return false; |
| 990 | } |
| 991 | |
| 992 | bool FillBlockDimLimit(GpuDeviceHandle device, BlockDim* block_dim_limit) { |
| 993 | // The BlockDim name is a mismatch against these GRID_DIM_* queries because |