| 42 | } |
| 43 | |
| 44 | TEST(RedzoneAllocatorTest, WriteToRedzone) { |
| 45 | constexpr int64 kRedzoneSize = 1 << 23; // 8MiB redzone on each side |
| 46 | // Redzone pattern should not be equal to zero; otherwise modify_redzone will |
| 47 | // break. |
| 48 | constexpr uint8 kRedzonePattern = 0x7e; |
| 49 | |
| 50 | // Allocate 32MiB + 1 byte (to make things misaligned) |
| 51 | constexpr int64 kAllocSize = (1 << 25) + 1; |
| 52 | |
| 53 | Platform* platform = |
| 54 | MultiPlatformManager::PlatformWithName("cuda").ValueOrDie(); |
| 55 | StreamExecutor* stream_exec = platform->ExecutorForDevice(0).ValueOrDie(); |
| 56 | cuda::PtxCompilationOptions opts; |
| 57 | StreamExecutorMemoryAllocator se_allocator(platform, {stream_exec}); |
| 58 | |
| 59 | Stream stream(stream_exec); |
| 60 | stream.Init(); |
| 61 | RedzoneAllocator allocator( |
| 62 | &stream, &se_allocator, opts, |
| 63 | /*memory_limit=*/RedzoneAllocator::kDefaultMemoryLimit, |
| 64 | /*redzone_size=*/kRedzoneSize, |
| 65 | /*redzone_pattern=*/kRedzonePattern); |
| 66 | TF_ASSERT_OK_AND_ASSIGN(DeviceMemory<uint8> buf, |
| 67 | allocator.AllocateBytes(/*byte_size=*/kAllocSize)); |
| 68 | EXPECT_REDZONE_OK(allocator.CheckRedzones()); |
| 69 | |
| 70 | char* buf_addr = reinterpret_cast<char*>(buf.opaque()); |
| 71 | DeviceMemoryBase lhs_redzone(buf_addr - kRedzoneSize, kRedzoneSize); |
| 72 | DeviceMemoryBase rhs_redzone(buf_addr + kAllocSize, kRedzoneSize); |
| 73 | |
| 74 | // Check that the redzones are in fact filled with kRedzonePattern. |
| 75 | auto check_redzone = [&](DeviceMemoryBase redzone, absl::string_view name) { |
| 76 | std::vector<uint8> host_buf(kRedzoneSize); |
| 77 | TF_ASSERT_OK(stream.ThenMemcpy(host_buf.data(), redzone, kRedzoneSize) |
| 78 | .BlockHostUntilDone()); |
| 79 | const int64 kMaxMismatches = 16; |
| 80 | int64 mismatches = 0; |
| 81 | for (int64 i = 0; i < host_buf.size(); ++i) { |
| 82 | if (mismatches == kMaxMismatches) { |
| 83 | ADD_FAILURE() << "Hit max number of mismatches; skipping others."; |
| 84 | break; |
| 85 | } |
| 86 | if (host_buf[i] != kRedzonePattern) { |
| 87 | ++mismatches; |
| 88 | EXPECT_EQ(host_buf[i], kRedzonePattern) |
| 89 | << "at index " << i << " of " << name << " redzone"; |
| 90 | } |
| 91 | } |
| 92 | }; |
| 93 | check_redzone(lhs_redzone, "lhs"); |
| 94 | check_redzone(rhs_redzone, "rhs"); |
| 95 | |
| 96 | // Modifies a redzone, checks that RedzonesAreUnmodified returns false, then |
| 97 | // reverts it back to its original value and checks that RedzonesAreUnmodified |
| 98 | // returns true. |
| 99 | auto modify_redzone = [&](DeviceMemoryBase redzone, int64 offset, |
| 100 | absl::string_view name) { |
| 101 | SCOPED_TRACE(absl::StrCat(name, ", offset=", offset)); |
nothing calls this directly
no test coverage detected