| /* Copyright 2019 The TensorFlow Authors. All Rights Reserved. |
| |
| Licensed under the Apache License, Version 2.0 (the "License"); |
| you may not use this file except in compliance with the License. |
| You may obtain a copy of the License at |
| |
| http://www.apache.org/licenses/LICENSE-2.0 |
| |
| Unless required by applicable law or agreed to in writing, software |
| distributed under the License is distributed on an "AS IS" BASIS, |
| WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| See the License for the specific language governing permissions and |
| limitations under the License. |
| ==============================================================================*/ |
| |
| #include "tensorflow/stream_executor/gpu/redzone_allocator.h" |
| |
| #include "absl/base/call_once.h" |
| #include "absl/container/fixed_array.h" |
| #include "absl/strings/str_format.h" |
| #include "absl/types/optional.h" |
| #include "tensorflow/core/framework/allocator.h" |
| #include "tensorflow/core/lib/core/errors.h" |
| #include "tensorflow/core/lib/core/status.h" |
| #include "tensorflow/stream_executor/device_memory.h" |
| #include "tensorflow/stream_executor/gpu/asm_compiler.h" |
| #include "tensorflow/stream_executor/kernel.h" |
| #include "tensorflow/stream_executor/kernel_spec.h" |
| #include "tensorflow/stream_executor/stream.h" |
| #include "tensorflow/stream_executor/stream_executor_pimpl.h" |
| |
| namespace stream_executor { |
| |
| // Rounds the value up to a multiple of the divisor by first calling CeilOfRatio |
| // then multiplying by the divisor. For example: RoundUpToNearest(13, 8) => 16 |
| template <typename T> |
| static T RoundUpToNearest(T value, T divisor) { |
| return tensorflow::MathUtil::CeilOfRatio(value, divisor) * divisor; |
| } |
| |
| // The size of the redzone at the end of the user buffer is rounded up to a |
| // multiple of kRhsRedzoneAlign. This simplifies the implementation a bit. |
| constexpr int64 kRhsRedzoneAlign = 4; |
| |
| using RedzoneCheckStatus = RedzoneAllocator::RedzoneCheckStatus; |
| |
| RedzoneAllocator::RedzoneAllocator(Stream* stream, |
| DeviceMemoryAllocator* memory_allocator, |
| GpuAsmOpts ptx_compilation_opts, |
| int64 memory_limit, int64 redzone_size, |
| uint8 redzone_pattern) |
| : device_ordinal_(stream->parent()->device_ordinal()), |
| stream_(stream), |
| memory_limit_(memory_limit), |
| redzone_size_(RoundUpToNearest( |
| redzone_size, |
| static_cast<int64>(tensorflow::Allocator::kAllocatorAlignment))), |
| redzone_pattern_(redzone_pattern), |
| memory_allocator_(memory_allocator), |
| gpu_compilation_opts_(ptx_compilation_opts) {} |
| |
| port::StatusOr<DeviceMemory<uint8>> RedzoneAllocator::AllocateBytes( |
| int64 byte_size) { |
| CHECK_GE(byte_size, 0) << "byte_size must be positive."; |
| if (byte_size > GetMemoryLimitInBytes()) { |
| return port::Status( |
| port::error::RESOURCE_EXHAUSTED, |
| absl::StrFormat( |
| "Allocating %d bytes exceeds the memory limit of %d bytes.", |
| byte_size, GetMemoryLimitInBytes())); |
| } |
| |
| int64 rhs_slop = RoundUpToNearest(byte_size, kRhsRedzoneAlign) - byte_size; |
| TF_ASSIGN_OR_RETURN( |
| OwningDeviceMemory allocated_buffer, |
| memory_allocator_->Allocate(device_ordinal_, |
| byte_size + 2 * redzone_size_ + rhs_slop, |
| /*retry_on_failure=*/false)); |
| allocated_bytes_excluding_redzones_ += byte_size; |
| |
| static_assert(sizeof(uint8) == 1, "Unexpected size"); |
| DeviceMemory<uint8> allocated_buffer_memory(*allocated_buffer); |
| |
| DeviceMemory<uint8> lhs_redzone = stream_->parent()->GetSubBuffer( |
| &allocated_buffer_memory, 0, redzone_size_); |
| |
| DeviceMemory<uint8> data_chunk = stream_->parent()->GetSubBuffer( |
| &allocated_buffer_memory, redzone_size_, byte_size); |
| |
| // Split up the RHS redzone into two pieces: |
| // - 0 to kRhsRedzoneAlign bytes adjacent to the user buffer, followed by |
| // - redzone_size_ bytes. |
| // We do this because Stream::ThenMemset32 requires the buffer address and |
| // size to be aligned to 4 bytes. |
| DeviceMemory<uint8> rhs_redzone_slop = stream_->parent()->GetSubBuffer( |
| &allocated_buffer_memory, redzone_size_ + byte_size, rhs_slop); |
| |
| DeviceMemory<uint8> rhs_redzone_nonslop = stream_->parent()->GetSubBuffer( |
| &allocated_buffer_memory, redzone_size_ + byte_size + rhs_slop, |
| redzone_size_); |
| |
| uint8 pattern_arr[] = {redzone_pattern_, redzone_pattern_, redzone_pattern_, |
| redzone_pattern_}; |
| uint32 pattern32; |
| std::memcpy(&pattern32, pattern_arr, sizeof(pattern32)); |
| stream_->ThenMemset32(&lhs_redzone, pattern32, redzone_size_); |
| if (rhs_slop != 0) { |
| stream_->ThenMemcpy(&rhs_redzone_slop, &pattern32, rhs_slop); |
| } |
| stream_->ThenMemset32(&rhs_redzone_nonslop, pattern32, redzone_size_); |
| |
| allocated_buffers_.emplace_back(std::move(allocated_buffer), byte_size); |
| return data_chunk; |
| } |
| |
| // PTX blob for the function which checks that every byte in |
| // input_buffer (length is buffer_length) is equal to redzone_pattern. |
| // |
| // On mismatch, increment the counter pointed to by out_mismatch_cnt_ptr. |
| // |
| // Generated from: |
| // __global__ void redzone_checker(unsigned char* input_buffer, |
| // unsigned char redzone_pattern, |
| // unsigned long long buffer_length, |
| // int* out_mismatched_ptr) { |
| // unsigned long long idx = threadIdx.x + blockIdx.x * blockDim.x; |
| // if (idx >= buffer_length) return; |
| // if (input_buffer[idx] != redzone_pattern) atomicAdd(out_mismatched_ptr, 1); |
| // } |
| // |
| // Code must compile for the oldest GPU XLA may be compiled for. |
| static const char* redzone_checker_ptx = R"( |
| .version 4.2 |
| .target sm_30 |
| .address_size 64 |
| |
| .visible .entry redzone_checker( |
| .param .u64 input_buffer, |
| .param .u8 redzone_pattern, |
| .param .u64 buffer_length, |
| .param .u64 out_mismatch_cnt_ptr |
| ) |
| { |
| .reg .pred %p<3>; |
| .reg .b16 %rs<3>; |
| .reg .b32 %r<6>; |
| .reg .b64 %rd<8>; |
| |
| ld.param.u64 %rd6, [buffer_length]; |
| mov.u32 %r1, %tid.x; |
| mov.u32 %r2, %ctaid.x; |
| mov.u32 %r3, %ntid.x; |
| mad.lo.s32 %r4, %r3, %r2, %r1; |
| cvt.u64.u32 %rd3, %r4; |
| setp.ge.u64 %p1, %rd3, %rd6; |
| @%p1 bra LBB6_3; |
| ld.param.u8 %rs1, [redzone_pattern]; |
| ld.param.u64 %rd4, [input_buffer]; |
| cvta.to.global.u64 %rd2, %rd4; |
| add.s64 %rd7, %rd2, %rd3; |
| ld.global.u8 %rs2, [%rd7]; |
| setp.eq.s16 %p2, %rs2, %rs1; |
| @%p2 bra LBB6_3; |
| ld.param.u64 %rd5, [out_mismatch_cnt_ptr]; |
| cvta.to.global.u64 %rd1, %rd5; |
| atom.global.add.u32 %r5, [%rd1], 1; |
| LBB6_3: |
| ret; |
| } |
| )"; |
| |
| // The PTX in redzone_checker_ptx has to be launched with specified types |
| // in the specified order. |
| using ComparisonKernelT = |
| TypedKernel<DeviceMemory<uint8>, uint8, uint64, DeviceMemory<uint64>>; |
| |
| // Check that redzones weren't overwritten on a host. |
| // |
| // Slower, but gives a more useful error message. |
| static port::StatusOr<RedzoneCheckStatus> CheckRedzoneHost( |
| DeviceMemoryBase redzone, DeviceMemoryBase user_allocation, |
| absl::string_view name, Stream* stream, uint8 redzone_pattern) { |
| uint64 size = redzone.size(); |
| auto redzone_data = absl::make_unique<uint8[]>(size); |
| TF_RETURN_IF_ERROR(stream->ThenMemcpy(redzone_data.get(), redzone, size) |
| .BlockHostUntilDone()); |
| |
| std::array<uint8, sizeof(uint64)> pattern_arr; |
| pattern_arr.fill(redzone_pattern); |
| uint64 pattern64; |
| std::memcpy(&pattern64, pattern_arr.data(), sizeof(uint64)); |
| |
| int64 i; |
| for (i = 0; i + 7 < size; i += sizeof(uint64)) { |
| uint64 rz_value = *reinterpret_cast<uint64*>(&redzone_data[i]); |
| if (rz_value != pattern64) { |
| return RedzoneCheckStatus(name, user_allocation.opaque(), i, pattern64, |
| rz_value); |
| } |
| } |
| for (; i < size; ++i) { |
| uint8 rz_value = redzone_data[i]; |
| if (rz_value != redzone_pattern) { |
| return RedzoneCheckStatus(name, user_allocation.opaque(), i, |
| redzone_pattern, rz_value); |
| } |
| } |
| return RedzoneCheckStatus::OK(); |
| } |
| |
| // Run the redzone checker on the provided buffer redzone. |
| // |
| // Increment out_param if mismatch occurs. |
| static void RunRedzoneChecker(Stream* stream, |
| const DeviceMemory<uint8>& redzone, |
| uint8 redzone_pattern, |
| const DeviceMemory<uint64>& out_param, |
| const ComparisonKernelT& comparison_kernel) { |
| StreamExecutor* executor = stream->parent(); |
| |
| int64 num_elements = redzone.size(); |
| int64 threads_per_block = std::min( |
| executor->GetDeviceDescription().threads_per_block_limit(), num_elements); |
| int64 block_count = |
| tensorflow::MathUtil::CeilOfRatio(num_elements, threads_per_block); |
| |
| stream->ThenLaunch(ThreadDim(threads_per_block), BlockDim(block_count), |
| comparison_kernel, redzone, redzone_pattern, |
| redzone.size(), out_param); |
| } |
| |
| // Since we reuse the same buffer for multiple checks, we re-initialize redzone |
| // with a NaN pattern after a failed check. |
| // |
| // This function is blocking, since redzone failing is a rare event. |
| static port::Status ReinitializeRedzone(Stream* stream, |
| DeviceMemoryBase redzone, |
| uint8 redzone_pattern) { |
| absl::FixedArray<uint8> redzone_array(redzone.size()); |
| redzone_array.fill(redzone_pattern); |
| stream->ThenMemcpy(&redzone, redzone_array.data(), redzone.size()); |
| TF_RETURN_IF_ERROR(stream->BlockHostUntilDone()); |
| return port::Status::OK(); |
| } |
| |
| // Check redzones around the user allocation. |
| // |
| // Precondition: the memory pointed out by out_param is zeroed. |
| static port::StatusOr<RedzoneCheckStatus> CheckRedzonesForBuffer( |
| Stream* stream, DeviceMemoryBase memory, |
| const DeviceMemory<uint64>& out_param, |
| const ComparisonKernelT& comparison_kernel, int64 user_allocation_size, |
| uint64 redzone_size, uint8 redzone_pattern) { |
| StreamExecutor* executor = stream->parent(); |
| int64 rhs_slop = |
| RoundUpToNearest<int64>(user_allocation_size, kRhsRedzoneAlign) - |
| user_allocation_size; |
| CHECK_EQ(memory.size(), user_allocation_size + rhs_slop + 2 * redzone_size); |
| |
| DeviceMemory<uint8> buffer_uint8(memory); |
| DeviceMemory<uint8> lhs_redzone = |
| executor->GetSubBuffer(&buffer_uint8, 0, |
| /*element_count=*/redzone_size); |
| DeviceMemory<uint8> user_allocation = |
| executor->GetSubBuffer(&buffer_uint8, redzone_size, |
| /*element_count=*/user_allocation_size); |
| DeviceMemory<uint8> rhs_redzone = |
| executor->GetSubBuffer(&buffer_uint8, redzone_size + user_allocation_size, |
| /*element_count=*/redzone_size + rhs_slop); |
| |
| RunRedzoneChecker(stream, lhs_redzone, redzone_pattern, out_param, |
| comparison_kernel); |
| RunRedzoneChecker(stream, rhs_redzone, redzone_pattern, out_param, |
| comparison_kernel); |
| int64 result; |
| CHECK_EQ(out_param.size(), sizeof(result)); |
| stream->ThenMemcpy(&result, out_param, sizeof(result)); |
| TF_RETURN_IF_ERROR(stream->BlockHostUntilDone()); |
| |
| if (result != 0) { |
| TF_ASSIGN_OR_RETURN(RedzoneCheckStatus lhs_check, |
| CheckRedzoneHost(lhs_redzone, user_allocation, "LHS", |
| stream, redzone_pattern)); |
| TF_ASSIGN_OR_RETURN(RedzoneCheckStatus rhs_check, |
| CheckRedzoneHost(rhs_redzone, user_allocation, "RHS", |
| stream, redzone_pattern)); |
| |
| CHECK(!lhs_check.ok() || !rhs_check.ok()) |
| << "Mismatched results with host and device comparison"; |
| |
| TF_RETURN_IF_ERROR( |
| ReinitializeRedzone(stream, lhs_redzone, redzone_pattern)); |
| TF_RETURN_IF_ERROR( |
| ReinitializeRedzone(stream, rhs_redzone, redzone_pattern)); |
| return !lhs_check.ok() ? lhs_check : rhs_check; |
| } |
| |
| return RedzoneCheckStatus::OK(); |
| } |
| |
| port::StatusOr<RedzoneCheckStatus> RedzoneAllocator::CheckRedzones() const { |
| StreamExecutor* executor = stream_->parent(); |
| |
| absl::Span<const uint8> compiled_ptx = {}; |
| port::StatusOr<absl::Span<const uint8>> compiled_ptx_or = |
| CompileGpuAsmOrGetCached(executor->device_ordinal(), redzone_checker_ptx, |
| gpu_compilation_opts_); |
| if (compiled_ptx_or.ok()) { |
| compiled_ptx = compiled_ptx_or.ValueOrDie(); |
| } else { |
| static absl::once_flag ptxas_not_found_logged; |
| absl::call_once(ptxas_not_found_logged, [&]() { |
| LOG(WARNING) << compiled_ptx_or.status().ToString() |
| << "\nRelying on driver to perform ptx compilation. " |
| << "\nModify $PATH to customize ptxas location." |
| << "\nThis message will be only logged once."; |
| }); |
| } |
| |
| ScopedDeviceMemory<uint64> out_param = |
| executor->AllocateOwnedScalar<uint64>(); |
| stream_->ThenMemZero(out_param.ptr(), sizeof(uint64)); |
| |
| TF_ASSIGN_OR_RETURN( |
| std::unique_ptr<ComparisonKernelT> comparison_kernel, |
| (executor->CreateTypedKernel<DeviceMemory<uint8>, uint8, uint64, |
| DeviceMemory<uint64>>( |
| "redzone_checker", redzone_checker_ptx, compiled_ptx))); |
| |
| for (const auto& buf_and_size : allocated_buffers_) { |
| TF_ASSIGN_OR_RETURN( |
| RedzoneCheckStatus redzone_status, |
| CheckRedzonesForBuffer(stream_, *buf_and_size.first, out_param.cref(), |
| *comparison_kernel, buf_and_size.second, |
| redzone_size_, redzone_pattern_)); |
| if (!redzone_status.ok()) { |
| return redzone_status; |
| } |
| } |
| |
| return RedzoneCheckStatus::OK(); |
| } |
| |
| std::string RedzoneCheckStatus::RedzoneFailureMsg() const { |
| return absl::StrFormat( |
| "Redzone mismatch in %s redzone of buffer %p at offset %d; " |
| "expected %08x but was %08x.", |
| buffer_name, user_buffer_address, offset, expected_value, actual_value); |
| } |
| |
| } // namespace stream_executor |
