| 298 | } |
| 299 | |
| 300 | port::StatusOr<RedzoneCheckStatus> RedzoneAllocator::CheckRedzones() const { |
| 301 | StreamExecutor* executor = stream_->parent(); |
| 302 | |
| 303 | absl::Span<const uint8> compiled_ptx = {}; |
| 304 | port::StatusOr<absl::Span<const uint8>> compiled_ptx_or = |
| 305 | cuda::CompilePtxOrGetCached(executor->device_ordinal(), |
| 306 | redzone_checker_ptx, ptx_compilation_opts_); |
| 307 | if (compiled_ptx_or.ok()) { |
| 308 | compiled_ptx = compiled_ptx_or.ValueOrDie(); |
| 309 | } else { |
| 310 | static std::once_flag ptxas_not_found_logged; |
| 311 | std::call_once(ptxas_not_found_logged, [&]() { |
| 312 | LOG(WARNING) << compiled_ptx_or.status().ToString() |
| 313 | << "\nRelying on driver to perform ptx compilation. " |
| 314 | << "This message will be only logged once."; |
| 315 | }); |
| 316 | } |
| 317 | |
| 318 | ScopedDeviceMemory<uint64> out_param = |
| 319 | executor->AllocateOwnedScalar<uint64>(); |
| 320 | stream_->ThenMemZero(out_param.ptr(), sizeof(uint64)); |
| 321 | |
| 322 | TF_ASSIGN_OR_RETURN( |
| 323 | std::shared_ptr<ComparisonKernelT> loaded_kernel, |
| 324 | (LoadKernelOrGetPtr<DeviceMemory<uint8>, uint8, uint64, |
| 325 | DeviceMemory<uint64>>( |
| 326 | executor, "redzone_checker", redzone_checker_ptx, compiled_ptx))); |
| 327 | |
| 328 | for (const auto& buf_and_size : allocated_buffers_) { |
| 329 | TF_ASSIGN_OR_RETURN( |
| 330 | RedzoneCheckStatus redzone_status, |
| 331 | CheckRedzonesForBuffer(stream_, *buf_and_size.first, out_param.cref(), |
| 332 | *loaded_kernel, buf_and_size.second, |
| 333 | redzone_size_, redzone_pattern_)); |
| 334 | if (!redzone_status.ok()) { |
| 335 | return redzone_status; |
| 336 | } |
| 337 | } |
| 338 | |
| 339 | return RedzoneCheckStatus::OK(); |
| 340 | } |
| 341 | |
| 342 | std::string RedzoneCheckStatus::RedzoneFailureMsg() const { |
| 343 | return absl::StrFormat( |