ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-04-19 22:39:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

limit the rotating count to prevent oom (#3087)

Browse Source
This commit is contained in:
Max Podkorytov
2025-10-24 08:55:34 -07:00
committed by GitHub
parent 775b96ea6a
commit fdcc1f75c3
1 changed files with 7 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
include/ck_tile/host/rotating_buffers.hpp
View File

@@ -29,12 +29,12 @@ struct RotatingMemWrapper
RotatingMemWrapper() = delete;
RotatingMemWrapper(const void* a_ptr_,
const void* b_ptr_,
std::size_t rotating_count_,
std::size_t rotating_count_hint,
std::size_t size_a_,
std::size_t size_b_)
: a_ptr(a_ptr_),
b_ptr(b_ptr_),
rotating_count(rotating_count_),
rotating_count(rotating_count_hint),
size_a(size_a_),
size_b(size_b_)
{
@@ -42,6 +42,11 @@ struct RotatingMemWrapper
p_a_grids.push_back(a_ptr);
p_b_grids.push_back(b_ptr);
// limit the rotating count to prevent oom
const uint64_t footprint = (size_a + size_b);
const uint64_t max_rotating_count = (1ULL << 31) / footprint;
rotating_count = std::min(rotating_count, max_rotating_count);
// Create (rotating_count - 1) additional copies at different memory addresses
for(size_t i = 1; i < rotating_count; i++)
{