1/*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements. See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership. The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License. You may obtain a copy of the License at
9 *
10 * http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied. See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19
20/*!
21 * \file src/relay/collage/priority_queue.h
22 * \brief An updatable priority queue.
23 */
24
25#ifndef TVM_RELAY_COLLAGE_PRIORITY_QUEUE_H_
26#define TVM_RELAY_COLLAGE_PRIORITY_QUEUE_H_
27
28#include <set>
29
30namespace tvm {
31namespace relay {
32namespace collage {
33
34/*! \brief Priority queue of search states, ordered by increasing cost. */
35template <typename T, typename CmpTPtr, typename EqTPtr>
36class PriorityQueue {
37 public:
38 PriorityQueue() = default;
39
40 /*! \brief Pushes \p item onto the queue. */
41 void Push(T* item) { set_.emplace(item); }
42
43 /*! \brief Pops the item with the least cost off the queue. */
44 T* Pop() {
45 ICHECK(!set_.empty());
46 T* item = *set_.begin();
47 set_.erase(set_.begin());
48 return item;
49 }
50
51 /*! \brief Updates the queue to account for \p item's best cost being lowered. */
52 void Update(T* item) {
53 auto itr = std::find_if(set_.begin(), set_.end(),
54 [item](const T* that) { return EqTPtr()(that, item); });
55 ICHECK(itr != set_.end());
56 set_.erase(itr);
57 set_.emplace(item);
58 }
59
60 bool empty() const { return set_.empty(); }
61 size_t size() const { return set_.size(); }
62
63 private:
64 // TODO(mbs): Actually use a pri-queue datastructure!
65 std::set<T*, CmpTPtr> set_;
66};
67
68} // namespace collage
69} // namespace relay
70} // namespace tvm
71
72#endif // TVM_RELAY_COLLAGE_PRIORITY_QUEUE_H_
73