已合并
host allocator support pinned_use_background_threads #28843
zhaoyu65创建于 2025年12月30日
host allocator support pinned_use_background_threads #28843
已合并
共 4 个文件变更+154-83
| @@ -0,0 +1,41 @@ | |||
| 1 | +import time | ||
| 2 | +import os | ||
| 3 | +import torch | ||
| 4 | +from torch.testing._internal.common_utils import TestCase, run_tests | ||
| 5 | +os.environ["PYTORCH_NPU_ALLOC_CONF"] = "pinned_use_background_threads:True" | ||
| 6 | + | ||
| 7 | +torch.manual_seed(0) | ||
| 8 | + | ||
| 9 | +DEVICE = "npu" | ||
| 10 | +ALLOC_SIZE = 1024 * 64 | ||
| 11 | +ITERS = 200 | ||
| 12 | + | ||
| 13 | + | ||
| 14 | +class TestPinnedMemoryBackgroundThreads(TestCase): | ||
| 15 | + | ||
| 16 | + def copy_tensor(copy_times=20): | ||
| 17 | + times = [] | ||
| 18 | + dummy = torch.empty(ALLOC_SIZE // 4, device=DEVICE) | ||
| 19 | + torch.npu.synchronize() | ||
| 20 | + streams = [torch.npu.Stream() for _ in range(8)] | ||
| 21 | + for i in range(copy_times): | ||
| 22 | + t0 = time.perf_counter_ns() | ||
| 23 | + buf = torch.empty( | ||
| 24 | + ALLOC_SIZE // 4, | ||
| 25 | + dtype=torch.float32, | ||
| 26 | + pin_memory=True | ||
| 27 | + ) | ||
| 28 | + | ||
| 29 | + with torch.npu.stream(streams[i % len(streams)]): | ||
| 30 | + dummy.copy_(buf, non_blocking=True) | ||
| 31 | + t1 = time.perf_counter_ns() | ||
| 32 | + times.append((t1 - t0) / 1e3) # us | ||
| 33 | + torch.npu.synchronize() | ||
| 34 | + return times | ||
| 35 | + | ||
| 36 | + | ||
| 37 | + def test_pinned_memory_background_threads(self): | ||
| 38 | + self.copy_tensor(ITERS) | ||
| 39 | + | ||
| 40 | +if __name__ == '__main__': | ||
| 41 | + run_tests() | ||
| @@ -184,6 +184,11 @@ public: | |||
| 184 | return getStats(); | 184 | return getStats(); |
| 185 | } | 185 | } |
| 186 | 186 | ||
| 187 | + bool pinned_use_background_threads() override | ||
| 188 | + { | ||
| 189 | + return c10_npu::NPUCachingAllocator::CachingAllocatorConfig::pinned_use_background_threads(); | ||
| 190 | + } | ||
| 191 | + | ||
| 187 | private: | 192 | private: |
| 188 | void allocate_host_memory(size_t size, void** ptr) override | 193 | void allocate_host_memory(size_t size, void** ptr) override |
| 189 | { | 194 | { |
| @@ -271,6 +271,19 @@ size_t CachingAllocatorConfig::roundup_power2_divisions(size_t size) | |||
| 271 | return instance().m_roundup_power2_divisions[index]; | 271 | return instance().m_roundup_power2_divisions[index]; |
| 272 | } | 272 | } |
| 273 | 273 | ||
| 274 | +size_t CachingAllocatorConfig::parsePinnedUseBackgroundThreads(const std::vector<std::string> &config, size_t i) | ||
| 275 | +{ | ||
| 276 | + consumeToken(config, ++i, ':'); | ||
| 277 | + if (++i < config.size()) { | ||
| 278 | + TORCH_CHECK(i < config.size() && (config[i] == "True" || config[i] == "False"), | ||
| 279 | + "Expected a single True/False argument for pinned_use_background_threads", PTA_ERROR(ErrCode::PARAM)); | ||
| 280 | + m_pinned_use_background_threads = (config[i] == "True"); | ||
| 281 | + } else { | ||
| 282 | + TORCH_CHECK(false, "Error, expecting pinned_use_background_threads value", PTA_ERROR(ErrCode::PARAM)); | ||
| 283 | + } | ||
| 284 | + return i; | ||
| 285 | +} | ||
| 286 | + | ||
| 274 | void CachingAllocatorConfig::parseArgs(const char *env, std::set<std::string> supported_settings) | 287 | void CachingAllocatorConfig::parseArgs(const char *env, std::set<std::string> supported_settings) |
| 275 | { | 288 | { |
| 276 | // If empty, set the default values | 289 | // If empty, set the default values |
| @@ -309,6 +322,8 @@ void CachingAllocatorConfig::parseArgs(const char *env, std::set<std::string> su | |||
| 309 | i = parseSegmentSizeMb(config, i); | 322 | i = parseSegmentSizeMb(config, i); |
| 310 | } else if (config[i] == "roundup_power2_divisions") { | 323 | } else if (config[i] == "roundup_power2_divisions") { |
| 311 | i = parseRoundUpPower2Divisions(config, i); | 324 | i = parseRoundUpPower2Divisions(config, i); |
| 325 | + } else if (config[i] == "pinned_use_background_threads") { | ||
| 326 | + i = parsePinnedUseBackgroundThreads(config, i); | ||
| 312 | } else { | 327 | } else { |
| 313 | TORCH_CHECK(false, "Unrecognized CachingAllocator option: ", config[i], PTA_ERROR(ErrCode::PARAM)); | 328 | TORCH_CHECK(false, "Unrecognized CachingAllocator option: ", config[i], PTA_ERROR(ErrCode::PARAM)); |
| 314 | } | 329 | } |
| @@ -4,117 +4,127 @@ | |||
| 4 | 4 | ||
| 5 | 5 | ||
| 6 | namespace c10_npu { | 6 | namespace c10_npu { |
| 7 | - namespace NPUCachingAllocator { | 7 | +namespace NPUCachingAllocator { |
| 8 | - constexpr size_t kAlignRoundLarge = 16384; // round up large allocs to 16 KB | 8 | +constexpr size_t kAlignRoundLarge = 16384; // round up large allocs to 16 KB |
| 9 | - constexpr size_t kSmallBuffer = 2097152; // "small" allocations are packed in 2 MiB blocks | 9 | +constexpr size_t kSmallBuffer = 2097152; // "small" allocations are packed in 2 MiB blocks |
| 10 | - constexpr size_t kLargeBuffer = 20971520; // "large" allocations may be packed in 20 MiB blocks | 10 | +constexpr size_t kLargeBuffer = 20971520; // "large" allocations may be packed in 20 MiB blocks |
| 11 | - constexpr size_t kMB = 1024 * 1024; // 1 MB | 11 | +constexpr size_t kMB = 1024 * 1024; // 1 MB |
| 12 | - constexpr size_t k20MB = 20; // 20 MB for segmemt_size | 12 | +constexpr size_t k20MB = 20; // 20 MB for segmemt_size |
| 13 | - constexpr size_t k512MB = 512; // 512 MB for segmemt_size | 13 | +constexpr size_t k512MB = 512; // 512 MB for segmemt_size |
| 14 | - constexpr size_t kRoundUpPowerOfTwoStart = 1ULL << 20; // 1 MB | 14 | +constexpr size_t kRoundUpPowerOfTwoStart = 1ULL << 20; // 1 MB |
| 15 | - constexpr size_t kRoundUpPowerOfTwoEnd = 1ULL << 36; // 64 GB | 15 | +constexpr size_t kRoundUpPowerOfTwoEnd = 1ULL << 36; // 64 GB |
| 16 | - constexpr size_t kRoundUpPowerOfTwoIntervals = 16; | 16 | +constexpr size_t kRoundUpPowerOfTwoIntervals = 16; |
| 17 | 17 | ||
| 18 | - class CachingAllocatorConfig { | 18 | +class CachingAllocatorConfig { |
| 19 | - public: | 19 | +public: |
| 20 | - static size_t max_split_size() | 20 | + static size_t max_split_size() |
| 21 | - { | 21 | + { |
| 22 | - return instance().m_max_split_size; | 22 | + return instance().m_max_split_size; |
| 23 | - } | 23 | + } |
| 24 | 24 | ||
| 25 | - static double garbage_collection_threshold() | 25 | + static double garbage_collection_threshold() |
| 26 | - { | 26 | + { |
| 27 | - return instance().m_garbage_collection_threshold; | 27 | + return instance().m_garbage_collection_threshold; |
| 28 | - } | 28 | + } |
| 29 | 29 | ||
| 30 | - static bool expandable_segments() | 30 | + static bool expandable_segments() |
| 31 | - { | 31 | + { |
| 32 | - return instance().m_expandable_segments; | 32 | + return instance().m_expandable_segments; |
| 33 | - } | 33 | + } |
| 34 | 34 | ||
| 35 | - static bool pin_memory_expandable_segments() | 35 | + static bool pin_memory_expandable_segments() |
| 36 | - { | 36 | + { |
| 37 | - return instance().m_pin_memory_expandable_segments; | 37 | + return instance().m_pin_memory_expandable_segments; |
| 38 | - } | 38 | + } |
| 39 | 39 | ||
| 40 | - static size_t base_addr_aligned_size() | 40 | + static size_t base_addr_aligned_size() |
| 41 | - { | 41 | + { |
| 42 | - return instance().m_base_addr_aligned_size; | 42 | + return instance().m_base_addr_aligned_size; |
| 43 | - } | 43 | + } |
| 44 | 44 | ||
| 45 | - static bool page_size_1g_enable() | 45 | + static bool page_size_1g_enable() |
| 46 | - { | 46 | + { |
| 47 | - return instance().m_page_size_1g; | 47 | + return instance().m_page_size_1g; |
| 48 | - } | 48 | + } |
| 49 | 49 | ||
| 50 | - static size_t segment_size_mb() | 50 | + static size_t segment_size_mb() |
| 51 | - { | 51 | + { |
| 52 | - return instance().m_segment_size_mb; | 52 | + return instance().m_segment_size_mb; |
| 53 | - } | 53 | + } |
| 54 | 54 | ||
| 55 | - static size_t roundup_power2_divisions(size_t size); | 55 | + static size_t roundup_power2_divisions(size_t size); |
| 56 | 56 | ||
| 57 | - static CachingAllocatorConfig &instance() | 57 | + static size_t pinned_use_background_threads() |
| 58 | - { | 58 | + { |
| 59 | - static CachingAllocatorConfig *s_instance = ([]() { | 59 | + return instance().m_pinned_use_background_threads; |
| 60 | - auto inst = new CachingAllocatorConfig(); | 60 | + } |
| 61 | - const char *env = getenv("PYTORCH_NPU_ALLOC_CONF"); | ||
| 62 | - inst->parseArgs(env); | ||
| 63 | - return inst; | ||
| 64 | - })(); | ||
| 65 | - return *s_instance; | ||
| 66 | - } | ||
| 67 | 61 | ||
| 68 | - void parseArgs(const char *env, std::set<std::string> supported_settings = {}); | 62 | + static CachingAllocatorConfig &instance() |
| 63 | + { | ||
| 64 | + static CachingAllocatorConfig *s_instance = ([]() { | ||
| 65 | + auto inst = new CachingAllocatorConfig(); | ||
| 66 | + const char *env = getenv("PYTORCH_NPU_ALLOC_CONF"); | ||
| 67 | + inst->parseArgs(env); | ||
| 68 | + return inst; | ||
| 69 | + })(); | ||
| 70 | + return *s_instance; | ||
| 71 | + } | ||
| 69 | 72 | ||
| 70 | - private: | 73 | + void parseArgs(const char *env, std::set<std::string> supported_settings = {}); |
| 71 | - size_t m_max_split_size; | ||
| 72 | 74 | ||
| 73 | - double m_garbage_collection_threshold; | 75 | +private: |
| 76 | + size_t m_max_split_size; | ||
| 74 | 77 | ||
| 75 | - bool m_expandable_segments; | 78 | + double m_garbage_collection_threshold; |
| 76 | 79 | ||
| 77 | - bool m_pin_memory_expandable_segments; | 80 | + bool m_expandable_segments; |
| 78 | 81 | ||
| 79 | - bool set_expandable_segments_flag = false; | 82 | + bool m_pin_memory_expandable_segments; |
| 80 | 83 | ||
| 81 | - size_t m_base_addr_aligned_size = kAlignRoundLarge; | 84 | + bool set_expandable_segments_flag = false; |
| 82 | 85 | ||
| 83 | - bool m_page_size_1g = false; // 新增1G页配置标志 | 86 | + size_t m_base_addr_aligned_size = kAlignRoundLarge; |
| 84 | 87 | ||
| 85 | - size_t m_segment_size_mb; | 88 | + bool m_page_size_1g = false; // 新增1G页配置标志 |
| 86 | 89 | ||
| 87 | - std::vector<size_t> m_roundup_power2_divisions; | 90 | + size_t m_segment_size_mb; |
| 88 | 91 | ||
| 89 | - CachingAllocatorConfig() | 92 | + std::vector<size_t> m_roundup_power2_divisions; |
| 90 | - : m_max_split_size(std::numeric_limits<size_t>::max()), | ||
| 91 | - m_garbage_collection_threshold(0), | ||
| 92 | - m_expandable_segments(false), | ||
| 93 | - m_pin_memory_expandable_segments(false), | ||
| 94 | - m_base_addr_aligned_size(kAlignRoundLarge), | ||
| 95 | - m_segment_size_mb(0), | ||
| 96 | - m_roundup_power2_divisions(kRoundUpPowerOfTwoIntervals, 0) | ||
| 97 | - {} | ||
| 98 | 93 | ||
| 99 | - void lexArgs(const char *env, std::vector<std::string> &config); | 94 | + std::atomic<bool> m_pinned_use_background_threads; // A flag to enable background thread for processing events. |
| 100 | 95 | ||
| 101 | - void consumeToken(const std::vector<std::string> &config, size_t i, const char c); | 96 | + CachingAllocatorConfig() |
| 97 | + : m_max_split_size(std::numeric_limits<size_t>::max()), | ||
| 98 | + m_garbage_collection_threshold(0), | ||
| 99 | + m_expandable_segments(false), | ||
| 100 | + m_pin_memory_expandable_segments(false), | ||
| 101 | + m_base_addr_aligned_size(kAlignRoundLarge), | ||
| 102 | + m_segment_size_mb(0), | ||
| 103 | + m_roundup_power2_divisions(kRoundUpPowerOfTwoIntervals, 0), | ||
| 104 | + m_pinned_use_background_threads(false) | ||
| 105 | + {} | ||
| 102 | 106 | ||
| 103 | - size_t parseMaxSplitSize(const std::vector<std::string> &config, size_t i); | 107 | + void lexArgs(const char *env, std::vector<std::string> &config); |
| 104 | 108 | ||
| 105 | - size_t parseGarbageCollectionThreshold(const std::vector<std::string> &config, size_t i); | 109 | + void consumeToken(const std::vector<std::string> &config, size_t i, const char c); |
| 106 | 110 | ||
| 107 | - size_t parseExpandableSegments(const std::vector<std::string> &config, size_t i); | 111 | + size_t parseMaxSplitSize(const std::vector<std::string> &config, size_t i); |
| 108 | 112 | ||
| 109 | - size_t parsePinMemoryExpandableSegments(const std::vector<std::string> &config, size_t i); | 113 | + size_t parseGarbageCollectionThreshold(const std::vector<std::string> &config, size_t i); |
| 110 | 114 | ||
| 111 | - size_t parseAddrAlignSize(const std::vector<std::string> &config, size_t i); | 115 | + size_t parseExpandableSegments(const std::vector<std::string> &config, size_t i); |
| 112 | 116 | ||
| 113 | - size_t parsePageSize(const std::vector<std::string> &config, size_t i); | 117 | + size_t parsePinMemoryExpandableSegments(const std::vector<std::string> &config, size_t i); |
| 114 | 118 | ||
| 115 | - size_t parseSegmentSizeMb(const std::vector<std::string> &config, size_t i); | 119 | + size_t parseAddrAlignSize(const std::vector<std::string> &config, size_t i); |
| 116 | 120 | ||
| 117 | - size_t parseRoundUpPower2Divisions(const std::vector<std::string> &config, size_t i); | 121 | + size_t parsePageSize(const std::vector<std::string> &config, size_t i); |
| 118 | - }; | 122 | + |
| 119 | - } // namespace NPUCachingAllocator | 123 | + size_t parseSegmentSizeMb(const std::vector<std::string> &config, size_t i); |
| 124 | + | ||
| 125 | + size_t parseRoundUpPower2Divisions(const std::vector<std::string> &config, size_t i); | ||
| 126 | + | ||
| 127 | + size_t parsePinnedUseBackgroundThreads(const std::vector<std::string> &config, size_t i); | ||
| 128 | +}; | ||
| 129 | +} // namespace NPUCachingAllocator | ||
| 120 | } // namespace c10_npu | 130 | } // namespace c10_npu |