d-smirnov commented on a change in pull request #9704: URL: https://github.com/apache/tvm/pull/9704#discussion_r770509396
########## File path: src/tir/usmp/algo/hill_climb.cc ########## @@ -0,0 +1,356 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +/*! + * \file tir/analysis/usmp/algo/greedy_by_size.cc + * \brief Implement greedy by size memory planning algorithm + */ +#include <tvm/arith/analyzer.h> +#include <tvm/runtime/device_api.h> +#include <tvm/tir/builtin.h> +#include <tvm/tir/function.h> +#include <tvm/tir/stmt_functor.h> +#include <tvm/tir/usmp/utils.h> + +#include <algorithm> +#include <numeric> +#include <sstream> + +namespace tvm { +namespace tir { +namespace usmp { +namespace algo { +/*! + * \brief Rounds up the offset to satisfy the alignement requirement + */ +static size_t round_up_to_byte_alignment(const size_t& non_aligned_byte_offset, + const int& byte_alignment) { + return ((non_aligned_byte_offset + byte_alignment - 1) / byte_alignment) * byte_alignment; +} + +/*! + * \brief A helper function check whether a offset is valid given the constraints + */ +static bool IsValidPlacement(const PoolInfo& candidate_pool, const size_t& next_offset, + const size_t& size_bytes) { + if (candidate_pool->size_hint_bytes == -1) { + // this means pool is not bounded + return true; + } + auto pool_size = static_cast<size_t>(candidate_pool->size_hint_bytes->value); + auto max_address = next_offset + size_bytes; + if (max_address <= pool_size) { + return true; + } + return false; +} + +/*! + * \brief Selects a pool for placement in the given set of ordered pool candidates + */ +static PoolInfo SelectPlacementPool( + const BufferInfo& buf_info, + const std::unordered_map<PoolInfo, size_t, ObjectPtrHash, ObjectPtrEqual>& pool_offsets) { + // Here the pool candidates are ordered when it is consumed by the algorithm. + // This could be from order the user has specified. However, schedulers are + // welcome to change the order for performance reasons. + for (const auto& pool_info : buf_info->pool_candidates) { + if (pool_offsets.count(pool_info)) { + return pool_info; + } + } + CHECK(false) << "TVM USMP Error: the space available in the provided pools exceeded when " + "trying to allocate the buffer : " + << buf_info << "\n. Please increase the size_hints for memory pools."; + return PoolInfo(); +} + +struct _ptr_hash { + template <typename T> + size_t operator()(const T& a) const { + return std::hash<T>()(a); + } +}; + +using alloc_map_t = std::unordered_map<const BufferInfoNode*, PoolAllocation, _ptr_hash>; + +static void sort_vector(std::vector<BufferInfo>* buffer_info_vec) { + std::sort(buffer_info_vec->begin(), buffer_info_vec->end(), + [](const BufferInfo& a, const BufferInfo& b) { + if (a->size_bytes->value == b->size_bytes->value) { + if (a->conflicts.size() == b->conflicts.size()) { + auto a_name_hash = std::hash<std::string>{}(a->name_hint->data); + auto b_name_hash = std::hash<std::string>{}(b->name_hint->data); + return a_name_hash > b_name_hash; + } else { + return a->conflicts.size() > b->conflicts.size(); + } + } + return a->size_bytes->value > b->size_bytes->value; + }); +} + +/* + * Modified version of greedy allocation from greedy_by_size.cc + */ +static void greedy(std::vector<BufferInfo>* buffer_info_vec, alloc_map_t* pool_allocations) { + for (const auto& buf_info : *buffer_info_vec) { + std::unordered_map<PoolInfo, size_t, ObjectPtrHash, ObjectPtrEqual> pool_offset_candidates; + for (const auto& pool_info : buf_info->pool_candidates) { + if (algo::IsValidPlacement(pool_info, 0, buf_info->size_bytes->value)) { + pool_offset_candidates[pool_info] = 0; + } + } + + std::vector<const BufferInfoNode*> buf_conf; + for (const auto& conflict_buf_info_obj : buf_info->conflicts) { + const BufferInfoNode* conflict_buf_info = conflict_buf_info_obj.as<BufferInfoNode>(); + if (pool_allocations->end() != pool_allocations->find(conflict_buf_info)) { + buf_conf.push_back(conflict_buf_info); + } + } + + // extra sorting for pool offsets + std::sort(buf_conf.begin(), buf_conf.end(), [&pool_allocations](const auto* a, const auto* b) { + return pool_allocations->operator[](a)->byte_offset->value < + pool_allocations->operator[](b)->byte_offset->value; + }); + + for (const auto* conflict_buf_info : buf_conf) { + size_t next_offset = 0; + auto pool_allocation = pool_allocations->operator[](conflict_buf_info); + next_offset = pool_allocation->byte_offset + conflict_buf_info->size_bytes; + next_offset = round_up_to_byte_alignment(next_offset, conflict_buf_info->alignment->value); + if (!pool_offset_candidates.count(pool_allocation->pool_info)) { + continue; + } + if (IsValidPlacement(pool_allocation->pool_info, next_offset, buf_info->size_bytes->value)) { + if (next_offset > pool_offset_candidates[pool_allocation->pool_info] && + pool_offset_candidates[pool_allocation->pool_info] + + static_cast<size_t>(buf_info->size_bytes) > + static_cast<size_t>(pool_allocation->byte_offset)) { + pool_offset_candidates[pool_allocation->pool_info] = next_offset; + } + } else { + pool_offset_candidates.erase(pool_allocation->pool_info); + } + } + auto selected_pool = algo::SelectPlacementPool(buf_info, pool_offset_candidates); + pool_allocations->operator[](buf_info.as<BufferInfoNode>()) = + PoolAllocation(selected_pool, Integer(pool_offset_candidates[selected_pool])); + } +} + +/* + * Finds highes allocated memory address for each pool + */ +static std::unordered_map<PoolInfo, size_t, ObjectPtrHash, ObjectPtrEqual> find_highest( + alloc_map_t* pool_allocations) { + std::unordered_map<PoolInfo, size_t, ObjectPtrHash, ObjectPtrEqual> max_pool_size; + for (const auto it : *pool_allocations) { + const BufferInfoNode* buf = it.first; + const PoolAllocation& pa = it.second; + size_t high_sz = pa->byte_offset + buf->size_bytes; + if (max_pool_size[pa->pool_info] <= high_sz) { + max_pool_size[pa->pool_info] = high_sz; + } + } + return max_pool_size; +} + +/* + * Simulated annealing / Hill climb + * + * Works by continiously invoking modified 'greedy-by-size' allocation + * assessing the result and introduce permutations which hopefully + * will led to more 'compact' memory allocation. + */ +Map<BufferInfo, PoolAllocation> HillClimb(const Array<BufferInfo>& buffer_info_arr, + const Integer& desired_bytes) { +// rand_r does not exist on Windows platform +#if defined(__linux__) || defined(__ANDROID__) + unsigned int _seedp = 0; +#define rnd_func() rand_r(&_seedp) +#else +#define rnd_func() rand() +#endif + + std::vector<BufferInfo> buffer_info_vec; + for (const auto& buffer_info : buffer_info_arr) { + ICHECK(buffer_info->pool_candidates.size()) + << "Cannot process buffer \"" << buffer_info->name_hint << "\" with no pool candidates"; + buffer_info_vec.push_back(std::move(buffer_info)); + } + + sort_vector(&buffer_info_vec); + + // populate positional index map + std::unordered_map<const BufferInfoNode*, int, _ptr_hash> _pos_map; + for (size_t index = 0; index < buffer_info_vec.size(); ++index) { + _pos_map[buffer_info_vec[index].as<BufferInfoNode>()] = index; + } + + // size_t first_attempt_size = 0; + size_t total_size = 0; + int attempts = 0; + // int successful_iteration = 0; + + int swap_i1 = -1; + int swap_i2 = -1; + size_t desired_bytes_ = desired_bytes; + constexpr auto _max_attempts = 500; + alloc_map_t rollback_pool_allocations; + alloc_map_t result_pool_allocations; + alloc_map_t pool_allocations; + + auto swap_buffers = [&buffer_info_vec, &_pos_map](int i1, int i2) { + if (i1 == i2) return; + auto b1 = buffer_info_vec[i1]; + auto b2 = buffer_info_vec[i2]; + buffer_info_vec[i1] = b2; + buffer_info_vec[i2] = b1; + + _pos_map[b1.as<BufferInfoNode>()] = i2; + _pos_map[b2.as<BufferInfoNode>()] = i1; + }; + + auto _pos = [&_pos_map](const auto* e) { + auto it = _pos_map.find(e); + if (it != _pos_map.end()) { + return it->second; + } + LOG(FATAL) << "not indexed"; + return -1; + }; + + for (; attempts < _max_attempts; ++attempts) { + rollback_pool_allocations = std::move(pool_allocations); + greedy(&buffer_info_vec, &pool_allocations); + + // estimate result buffers + auto max_pool_size = find_highest(&pool_allocations); + + // calculate summary + size_t total = 0; + for (const auto& el : max_pool_size) { + total += el.second; + } + // accept/reject result heuristic + if (!total_size || + (total_size > total || + rnd_func() % 100 < static_cast<int>(300 * (total - total_size) / total / attempts))) { + // remember winning combination + result_pool_allocations = pool_allocations; + total_size = total; + + // reached desired size + if (total_size <= desired_bytes_) { + break; + } + + } else { + // rollback + swap_buffers(swap_i2, swap_i1); + pool_allocations = std::move(rollback_pool_allocations); + max_pool_size = find_highest(&pool_allocations); + } + + std::vector<const BufferInfoNode*> max_pool_buf; + + for (const auto& it : pool_allocations) { + const auto* buf = it.first; + const auto pa = it.second; + size_t high_sz = pa->byte_offset + buf->size_bytes; + if (max_pool_size[pa->pool_info] == high_sz) { + max_pool_buf.push_back(buf); + } + } + + // pick highest + const BufferInfoNode* suspect = max_pool_buf[rand() % max_pool_buf.size()]; + PoolAllocation suspect_pa = pool_allocations[suspect]; + + std::unordered_map<int, const BufferInfoNode*, _ptr_hash> first_level_set; Review comment: Added short description for levels -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@tvm.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
