1 /* Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_COMPILER_XLA_STREAM_EXECUTOR_GPU_REDZONE_ALLOCATOR_H_ 17 #define TENSORFLOW_COMPILER_XLA_STREAM_EXECUTOR_GPU_REDZONE_ALLOCATOR_H_ 18 19 #include <vector> 20 21 #include "tensorflow/compiler/xla/stream_executor/device_memory_allocator.h" 22 #include "tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.h" 23 #include "tensorflow/compiler/xla/stream_executor/gpu/gpu_asm_opts.h" 24 #include "tensorflow/core/lib/math/math_util.h" 25 #include "tensorflow/core/platform/stream_executor_no_cuda.h" 26 27 namespace stream_executor { 28 29 // An allocator that allocates a bit of extra memory around the beginning/end of 30 // every allocation and can check that this memory is unmodified. 31 // 32 // This can be used to check for out-of-bounds writes, and, if the redzone is 33 // filled with a sufficiently "ugly" pattern, may also be able to check for 34 // out-of-bounds reads. The default fill pattern of -1 is an unusual NaN 35 // pattern when interpreted as a floating-point number, so hopefully works for 36 // out-of-bounds reads and writes in those cases. 37 // 38 // This class implements ScratchAllocator, so can be used to allocate temp 39 // memory for cudnn convolutions. 40 class RedzoneAllocator : public ScratchAllocator { 41 public: 42 static constexpr int64_t kDefaultRedzoneSize = 43 1LL << 23; // 8MiB per side, 16MiB total. 44 static constexpr uint8 kDefaultRedzonePattern = -1; 45 RedzoneAllocator(Stream* stream, DeviceMemoryAllocator* memory_allocator, 46 GpuAsmOpts gpu_compilation_opts_, 47 int64_t memory_limit = (1LL << 32), // 4GB 48 int64_t redzone_size = kDefaultRedzoneSize, 49 uint8 redzone_pattern = kDefaultRedzonePattern); 50 51 // Redzones don't count towards the memory limit. GetMemoryLimitInBytes()52 int64_t GetMemoryLimitInBytes() override { return memory_limit_; } 53 TotalAllocatedBytesExcludingRedzones()54 int64_t TotalAllocatedBytesExcludingRedzones() const { 55 return allocated_bytes_excluding_redzones_; 56 } 57 58 port::StatusOr<DeviceMemory<uint8>> AllocateBytes(int64_t byte_size) override; 59 60 // Non-empty redzone check status implies that there was a write into a 61 // redzone, with a string communicating the location of the write. 62 struct RedzoneCheckStatus { 63 RedzoneCheckStatus() = default; 64 RedzoneCheckStatusRedzoneCheckStatus65 RedzoneCheckStatus(absl::string_view buffer_name, void* user_buffer_address, 66 int64_t offset, uint64_t expected_value, 67 uint64_t actual_value) 68 : buffer_name(buffer_name), 69 user_buffer_address(user_buffer_address), 70 offset(offset), 71 expected_value(expected_value), 72 actual_value(actual_value) {} 73 OKRedzoneCheckStatus74 static RedzoneCheckStatus OK() { return {}; } 75 okRedzoneCheckStatus76 bool ok() { return user_buffer_address == nullptr; } 77 78 std::string RedzoneFailureMsg() const; 79 80 std::string buffer_name = {}; 81 void* user_buffer_address = nullptr; 82 int64_t offset = 0; 83 uint64_t expected_value = 0; 84 uint64_t actual_value = 0; 85 }; 86 87 // Determines whether redzones around all allocated buffers are unmodified. 88 // 89 // Reinitializes redzones to the expected value, so that the same buffer 90 // could be reused for multiple checks. 91 // 92 // Returns: 93 // 94 // - RedzoneCheckStatus::OK() if everything went well. 95 // - RedzoneCheckStatus with a non-empty error message iff a write into a 96 // redzone has been detected. 97 // - A stream error, if loading or launching the kernel has failed. 98 port::StatusOr<RedzoneCheckStatus> CheckRedzones() const; 99 stream()100 Stream* stream() const { return stream_; } 101 102 private: 103 const int device_ordinal_; 104 Stream* stream_; 105 106 // Memory limit of the allocator in bytes. 107 const int64_t memory_limit_; 108 109 // Redzone size on *one side* of allocation in bytes. 110 // 111 // Must be a multiple of kXlaAllocatedBufferAlignBytes, otherwise the buffers 112 // returned to users will be misaligned. 113 const int64_t redzone_size_; 114 115 const uint8 redzone_pattern_; 116 DeviceMemoryAllocator* memory_allocator_; 117 GpuAsmOpts gpu_compilation_opts_; 118 119 // The second element of the pair is the size of the user allocation. This 120 // isn't necessarily just first.size() - 2 * redzone_size_ because when the 121 // user allocation size is not a multiple of 4 bytes, we round up the size of 122 // the RHS redzone. 123 // 124 // ScratchAllocators need to free all allocated memory on destruction so we 125 // use `OwningDeviceMemory` here. 126 std::vector<std::pair<OwningDeviceMemory, int64_t>> allocated_buffers_; 127 128 int64_t allocated_bytes_excluding_redzones_ = 0; 129 }; 130 131 } // namespace stream_executor 132 133 #endif // TENSORFLOW_COMPILER_XLA_STREAM_EXECUTOR_GPU_REDZONE_ALLOCATOR_H_ 134