Halide 22.0.0 Halide compiler and libraries

State.h
Go to the documentation of this file.
1#ifndef STATE_H
2#define STATE_H
3
4#include "ASLog.h"
5#include "CostModel.h"
6#include "DefaultCostModel.h"
7#include "Featurization.h"
8#include "FunctionDAG.h"
9#include "LoopNest.h"
10#include "PerfectHashMap.h"
11#include <set>
12#include <unordered_set>
13#include <vector>
14
15namespace Halide {
16namespace Internal {
17namespace Autoscheduler {
18
19using std::map;
20using std::pair;
21using std::set;
22using std::string;
23using std::unordered_set;
24using std::vector;
25
26 bool verify_memoized_features();
27
28 bool is_memoize_blocks_enabled();
29
30 double get_stack_memory_adjustment_factor();
31
32 constexpr int kLocalMemoryLimit = 524288; // 512 KB
33
34// Stack memory limit = Total GPU Memory / (# of SMs * maximum threads per SM)
35// = 103232 bytes
36// Not all 103232 bytes will be free for allocations so reduce it by factor to
37// allow a buffer
38 int64_t get_stack_memory_limit();
39
40 bool use_adjusted_tilings();
41
42 bool compute_root_and_inline_only();
43
44 struct NoOpMutator {
45 void operator()(LoopNest *new_loop_nest) const {
46 }
47};
48
49template<typename PostCreateMutator>
50 void deep_copy_loop_nest(LoopNest *new_loop_nest,
51 const LoopNest *new_loop_nest_parent,
52 const IntrusivePtr<const LoopNest> &existing_loop_nest,
53 const PostCreateMutator &post_create_mutator) {
54 new_loop_nest->copy_from(*existing_loop_nest);
55
56 for (std::size_t i = 0, N = new_loop_nest->children.size(); i < N; ++i) {
57 LoopNest *new_child = new LoopNest;
58 new_loop_nest->children[i] = new_child;
59 deep_copy_loop_nest(new_child, new_loop_nest, existing_loop_nest->children[i], post_create_mutator);
60 }
61
62 post_create_mutator(new_loop_nest);
63}
64
65 using LoopNestMap = map<const LoopNest *, pair<const LoopNest *, int>>;
66
67template<typename PostCreateMutator>
68 LoopNest *deep_copy_loop_nest(const IntrusivePtr<const LoopNest> &loop_nest,
69 const PostCreateMutator &post_create_mutator) {
70 LoopNest *new_loop_nest = new LoopNest;
71 deep_copy_loop_nest(new_loop_nest, nullptr, loop_nest, post_create_mutator);
72 return new_loop_nest;
73}
74
75struct State {
76 mutable RefCount ref_count;
77 IntrusivePtr<const LoopNest> root;
78 IntrusivePtr<const State> parent;
79 double cost = 0;
80 std::vector<double> cost_per_stage;
81 NodeMap<bool> always_consider_inline;
82 int num_decisions_made = 0;
83 bool penalized = false;
84 string schedule_source;
85
86 State() = default;
87 State(const State &) = delete;
88 State(State &&) = delete;
89 void operator=(const State &) = delete;
90 void operator=(State &&) = delete;
91
92 uint64_t structural_hash(int depth) const;
93
94 // Compute the parent and depth of every loop nest node
95 void compute_loop_nest_parents(LoopNestMap &p,
96 const LoopNest *here,
97 int depth) const;
98
99 const LoopNest *deepest_common_ancestor(const LoopNestMap &parent,
100 const LoopNest *a,
101 const LoopNest *b) const;
102
103 // We use the post_create_mutator so that the loop nests can be modified
104 // before they become IntrusivePtr<const LoopNest> as children and cannot be modified
105 template<typename PostCreateMutator>
106 LoopNest *create_feature_root(const PostCreateMutator &post_create_mutator) const {
107 LoopNest *new_root = new LoopNest;
108 deep_copy_loop_nest<PostCreateMutator>(new_root, nullptr, root, post_create_mutator);
109 return new_root;
110 }
111
112 bool has_loop_nest_without_thread_loops() const;
113
114 bool has_compute_root_loops_without_blocks() const;
115
116 struct FeatureLoopNestMutator {
117 const Anderson2021Params &params;
118 const Target &target;
119
120 void operator()(LoopNest *new_loop_nest) const;
121
122 // In phase 2, any compute_root loop marked 'none' will be split into
123 // blocks, threads, and serial loops. To enable the cost model to make a
124 // meaningful prediction on these pre-split loops, we assume a split into
125 // blocks and threads with a single full warp (if possible)
126 void split_compute_root_loops(LoopNest *loop_nest) const;
127
128 // If a loop nest does not have thread loops, split the outermost serial
129 // loops to create thread loops with extents 1
130 void add_outer_thread_loops(LoopNest *loop_nest) const;
131 };
132
133 IntrusivePtr<const LoopNest> get_root_for_features(const Anderson2021Params &params,
134 const Target &target) const;
135
136 void set_gpu_store_site(const LoopNestMap &parent,
137 const LoopNest *loop,
138 LoopNest::Sites &site) const;
139
140 bool compute_featurization(const FunctionDAG &dag,
141 const Anderson2021Params &params,
142 const Target &target,
143 StageMap<ScheduleFeatures> *features,
144 Statistics &stats,
145 bool verbose = false) const;
146
147 void save_featurization(const FunctionDAG &dag,
148 const Anderson2021Params &params,
149 const Target &target,
150 std::ostream &out) const;
151
152 bool contains_store_at(const set<const FunctionDAG::Node *> &outermost_store_at,
153 const IntrusivePtr<const LoopNest> &parent) const;
154
155 // For GPU, only allow store_at root or inside the outermost loop nest. Any
156 // store_ats further in will be hoisted and expanded, increasing the
157 // amount of shared memory required.
158 bool contains_store_at_further_in_than_outermost() const;
159
160 bool has_dynamic_allocation_inside_thread() const;
161
162 bool exceeds_serial_extents_limit(const Target &target) const;
163
164 int64_t get_shared_mem_alloc_size(const LoopNest *block,
165 const LoopNest *loop) const;
166
167 bool exceeds_shared_memory_limit(const Anderson2021Params &params,
168 const Target &target) const;
169
170 bool exceeds_local_memory_limit(const Anderson2021Params &params,
171 const Target &target) const;
172
173 bool calculate_cost(const FunctionDAG &dag,
174 const Anderson2021Params &params,
175 const Target &target,
176 CostModel *cost_model,
177 Statistics &stats,
178 bool verbose = false);
179
180 // Make a child copy of this state. The loop nest is const (we
181 // make mutated copies of it, rather than mutating it), so we can
182 // continue to point to the same one and so this is a cheap
183 // operation.
184 IntrusivePtr<State> make_child() const;
185
186 void dump() const;
187
188 void print_compute_locations() const;
189
190 void fuse_gpu_blocks(LoopNest::StageScheduleState *state,
191 Stage &stage,
192 const vector<VarOrRVar> &parallel_vars,
193 const vector<int64_t> &parallel_extents,
194 const vector<int> &constant_extents) const;
195
196 void mark_gpu_blocks(LoopNest::StageScheduleState *state,
197 Stage &stage,
198 const vector<VarOrRVar> &parallel_vars,
199 const vector<int64_t> &parallel_extents) const;
200
201 bool mark_gpu_threads(LoopNest::StageScheduleState *state,
202 Stage &stage,
203 std::unordered_set<std::string> &new_serial_vars,
204 std::ostringstream &staged_funcs_schedule_source) const;
205
206 bool can_fuse_gpu(const vector<int64_t> &parallel_extents) const;
207
208 // Apply the schedule represented by this state to a Halide
209 // Pipeline. Also generate source code for the schedule for the
210 // user to copy-paste to freeze this schedule as permanent artifact.
211 void apply_schedule(const FunctionDAG &dag,
212 const Anderson2021Params &params,
213 const Target &target);
214
215 bool should_always_consider_inline(const FunctionDAG::Node *node) const;
216 void add_to_always_consider_inline_options(const FunctionDAG::Node *node);
217 void update_always_consider_inline_options(const FunctionDAG::Node *node);
218
219 const LoopNest *deepest_valid_compute_location(const Anderson2021Params &params,
220 const LoopNestMap &parent,
221 const FunctionDAG::Node &node,
222 const LoopNest *loop,
223 const LoopNest *root,
224 StageMap<int64_t> &total_shared_mem_alloc_sizes) const;
225 int64_t total_loop_extents_of_ancestors(const LoopNestMap &parent,
226 const LoopNest *loop) const;
227};
228
229// A priority queue of states, sorted according to increasing
230// cost. Never shrinks, to avoid reallocations.
231// Can't use std::priority_queue because it doesn't support unique_ptr.
232 class StateQueue {
233private:
234 struct CompareStates {
235 bool operator()(const IntrusivePtr<State> &a, const IntrusivePtr<State> &b) const {
236 return a->cost > b->cost;
237 }
238 };
239
240 std::vector<IntrusivePtr<State>> storage;
241 size_t sz = 0;
242
243public:
244 void emplace(IntrusivePtr<State> &&s) {
245 if (sz >= storage.size()) {
246 storage.resize(std::max(sz * 2, (size_t)64));
247 }
248 internal_assert(sz < storage.size()) << sz << " " << storage.size() << "\n";
249 storage[sz] = std::move(s);
250 sz++;
251 std::push_heap(storage.begin(), storage.begin() + sz, CompareStates{});
252 }
253
254 IntrusivePtr<State> pop() {
255 internal_assert(sz <= storage.size()) << sz << " " << storage.size() << "\n";
256 std::pop_heap(storage.begin(), storage.begin() + sz, CompareStates{});
257 sz--;
258 return std::move(storage[sz]);
259 }
260
261 const IntrusivePtr<State> &top() {
262 return storage[0];
263 }
264
265 bool empty() const {
266 return sz == 0;
267 }
268
269 size_t size() const {
270 return sz;
271 }
272
273 void swap(StateQueue &other) noexcept {
274 storage.swap(other.storage);
275 std::swap(sz, other.sz);
276 }
277
278 IntrusivePtr<State> operator[](int idx) const {
279 return storage[idx];
280 }
281
282 void resort() {
283 std::make_heap(storage.begin(), storage.begin() + sz, CompareStates{});
284 }
285
286 void clear() {
287 for (size_t i = 0; i < sz; i++) {
288 storage[i] = IntrusivePtr<State>{};
289 }
290 sz = 0;
291 }
292};
293
294} // namespace Autoscheduler
295} // namespace Internal
296} // namespace Halide
297
298#endif // STATE_H
ASLog.h
internal_assert
#define internal_assert(c)
Definition: Error.h:232
PerfectHashMap.h
CostModel.h
DefaultCostModel.h
Featurization.h
FunctionDAG.h
LoopNest.h
Halide::CostModel
Definition: CostModel.h:61
Halide::Internal::Autoscheduler::StateQueue
Definition: State.h:232
Halide::Internal::Autoscheduler::StateQueue::top
const IntrusivePtr< State > & top()
Definition: State.h:261
Halide::Internal::Autoscheduler::StateQueue::clear
void clear()
Definition: State.h:286
Halide::Internal::Autoscheduler::StateQueue::empty
bool empty() const
Definition: State.h:265
Halide::Internal::Autoscheduler::StateQueue::emplace
void emplace(IntrusivePtr< State > &&s)
Definition: State.h:244
Halide::Internal::Autoscheduler::StateQueue::pop
IntrusivePtr< State > pop()
Definition: State.h:254
Halide::Internal::Autoscheduler::StateQueue::size
size_t size() const
Definition: State.h:269
Halide::Internal::Autoscheduler::StateQueue::swap
void swap(StateQueue &other) noexcept
Definition: State.h:273
Halide::Internal::Autoscheduler::StateQueue::operator[]
IntrusivePtr< State > operator[](int idx) const
Definition: State.h:278
Halide::Internal::Autoscheduler::StateQueue::resort
void resort()
Definition: State.h:282
Halide::Internal::RefCount
A class representing a reference count to be used with IntrusivePtr.
Definition: IntrusivePtr.h:19
Halide::Stage
A single definition of a Func.
Definition: Func.h:70
PerfectHashMap
Definition: PerfectHashMap.h:38
Halide::Internal::Autoscheduler::compute_root_and_inline_only
bool compute_root_and_inline_only()
Halide::Internal::Autoscheduler::get_stack_memory_limit
int64_t get_stack_memory_limit()
Halide::Internal::Autoscheduler::use_adjusted_tilings
bool use_adjusted_tilings()
Halide::Internal::Autoscheduler::LoopNestMap
map< const LoopNest *, pair< const LoopNest *, int > > LoopNestMap
Definition: State.h:65
Halide::Internal::Autoscheduler::kLocalMemoryLimit
constexpr int kLocalMemoryLimit
Definition: State.h:32
Halide::Internal::Autoscheduler::get_stack_memory_adjustment_factor
double get_stack_memory_adjustment_factor()
Halide::Internal::Autoscheduler::verify_memoized_features
bool verify_memoized_features()
Halide::Internal::Autoscheduler::is_memoize_blocks_enabled
bool is_memoize_blocks_enabled()
Halide::Internal::Autoscheduler::deep_copy_loop_nest
void deep_copy_loop_nest(LoopNest *new_loop_nest, const LoopNest *new_loop_nest_parent, const IntrusivePtr< const LoopNest > &existing_loop_nest, const PostCreateMutator &post_create_mutator)
Definition: State.h:50
Halide::Internal::max
ConstantInterval max(const ConstantInterval &a, const ConstantInterval &b)
Halide
This file defines the class FunctionDAG, which is our representation of a Halide pipeline,...
Definition: AbstractGenerator.h:19
Halide::LinkageType::Internal
@ Internal
Not visible externally, similar to 'static' linkage in C.
uint64_t
unsigned __INT64_TYPE__ uint64_t
Definition: runtime_internal.h:23
int64_t
signed __INT64_TYPE__ int64_t
Definition: runtime_internal.h:22
size_t
__SIZE_TYPE__ size_t
Definition: runtime_internal.h:31
Halide::Internal::Autoscheduler::Anderson2021Params
Definition: CostModel.h:18
Halide::Internal::Autoscheduler::FunctionDAG::Node
Definition: FunctionDAG.h:378
Halide::Internal::Autoscheduler::FunctionDAG
Definition: FunctionDAG.h:367
Halide::Internal::Autoscheduler::LoopNest::Sites
Definition: LoopNest.h:100
Halide::Internal::Autoscheduler::LoopNest::StageScheduleState
Definition: LoopNest.h:209
Halide::Internal::Autoscheduler::LoopNest
Definition: LoopNest.h:34
Halide::Internal::Autoscheduler::LoopNest::children
std::vector< IntrusivePtr< const LoopNest > > children
Definition: LoopNest.h:42
Halide::Internal::Autoscheduler::LoopNest::copy_from
void copy_from(const LoopNest &n)
Halide::Internal::Autoscheduler::NoOpMutator
Definition: State.h:44
Halide::Internal::Autoscheduler::NoOpMutator::operator()
void operator()(LoopNest *new_loop_nest) const
Definition: State.h:45
Halide::Internal::Autoscheduler::State::FeatureLoopNestMutator
Definition: State.h:116
Halide::Internal::Autoscheduler::State::FeatureLoopNestMutator::target
const Target & target
Definition: State.h:118
Halide::Internal::Autoscheduler::State::FeatureLoopNestMutator::operator()
void operator()(LoopNest *new_loop_nest) const
Halide::Internal::Autoscheduler::State::FeatureLoopNestMutator::split_compute_root_loops
void split_compute_root_loops(LoopNest *loop_nest) const
Halide::Internal::Autoscheduler::State::FeatureLoopNestMutator::add_outer_thread_loops
void add_outer_thread_loops(LoopNest *loop_nest) const
Halide::Internal::Autoscheduler::State::FeatureLoopNestMutator::params
const Anderson2021Params & params
Definition: State.h:117
Halide::Internal::Autoscheduler::State
Definition: State.h:21
Halide::Internal::Autoscheduler::State::penalized
bool penalized
Definition: State.h:32
Halide::Internal::Autoscheduler::State::save_featurization
void save_featurization(const FunctionDAG &dag, const Anderson2021Params &params, const Target &target, std::ostream &out) const
Halide::Internal::Autoscheduler::State::exceeds_serial_extents_limit
bool exceeds_serial_extents_limit(const Target &target) const
Halide::Internal::Autoscheduler::State::get_shared_mem_alloc_size
int64_t get_shared_mem_alloc_size(const LoopNest *block, const LoopNest *loop) const
Halide::Internal::Autoscheduler::State::has_compute_root_loops_without_blocks
bool has_compute_root_loops_without_blocks() const
Halide::Internal::Autoscheduler::State::deepest_valid_compute_location
const LoopNest * deepest_valid_compute_location(const Anderson2021Params &params, const LoopNestMap &parent, const FunctionDAG::Node &node, const LoopNest *loop, const LoopNest *root, StageMap< int64_t > &total_shared_mem_alloc_sizes) const
Halide::Internal::Autoscheduler::State::dump
void dump() const
Halide::Internal::Autoscheduler::State::update_always_consider_inline_options
void update_always_consider_inline_options(const FunctionDAG::Node *node)
Halide::Internal::Autoscheduler::State::total_loop_extents_of_ancestors
int64_t total_loop_extents_of_ancestors(const LoopNestMap &parent, const LoopNest *loop) const
Halide::Internal::Autoscheduler::State::can_fuse_gpu
bool can_fuse_gpu(const vector< int64_t > &parallel_extents) const
Halide::Internal::Autoscheduler::State::should_always_consider_inline
bool should_always_consider_inline(const FunctionDAG::Node *node) const
Halide::Internal::Autoscheduler::State::contains_store_at_further_in_than_outermost
bool contains_store_at_further_in_than_outermost() const
Halide::Internal::Autoscheduler::State::operator=
void operator=(const State &)=delete
Halide::Internal::Autoscheduler::State::compute_loop_nest_parents
void compute_loop_nest_parents(LoopNestMap &p, const LoopNest *here, int depth) const
Halide::Internal::Autoscheduler::State::State
State(const State &)=delete
Halide::Internal::Autoscheduler::State::has_dynamic_allocation_inside_thread
bool has_dynamic_allocation_inside_thread() const
Halide::Internal::Autoscheduler::State::fuse_gpu_blocks
void fuse_gpu_blocks(LoopNest::StageScheduleState *state, Stage &stage, const vector< VarOrRVar > &parallel_vars, const vector< int64_t > &parallel_extents, const vector< int > &constant_extents) const
Halide::Internal::Autoscheduler::State::num_decisions_made
int num_decisions_made
Definition: State.h:30
Halide::Internal::Autoscheduler::State::set_gpu_store_site
void set_gpu_store_site(const LoopNestMap &parent, const LoopNest *loop, LoopNest::Sites &site) const
Halide::Internal::Autoscheduler::State::print_compute_locations
void print_compute_locations() const
Halide::Internal::Autoscheduler::State::mark_gpu_threads
bool mark_gpu_threads(LoopNest::StageScheduleState *state, Stage &stage, std::unordered_set< std::string > &new_serial_vars, std::ostringstream &staged_funcs_schedule_source) const
Halide::Internal::Autoscheduler::State::always_consider_inline
NodeMap< bool > always_consider_inline
Definition: State.h:81
Halide::Internal::Autoscheduler::State::parent
IntrusivePtr< const State > parent
Definition: State.h:26
Halide::Internal::Autoscheduler::State::operator=
void operator=(State &&)=delete
Halide::Internal::Autoscheduler::State::cost
double cost
Definition: State.h:28
Halide::Internal::Autoscheduler::State::schedule_source
string schedule_source
Definition: State.h:36
Halide::Internal::Autoscheduler::State::structural_hash
uint64_t structural_hash(int depth) const
Halide::Internal::Autoscheduler::State::root
IntrusivePtr< const LoopNest > root
Definition: State.h:24
Halide::Internal::Autoscheduler::State::exceeds_shared_memory_limit
bool exceeds_shared_memory_limit(const Anderson2021Params &params, const Target &target) const
Halide::Internal::Autoscheduler::State::compute_featurization
bool compute_featurization(const FunctionDAG &dag, const Anderson2021Params &params, const Target &target, StageMap< ScheduleFeatures > *features, Statistics &stats, bool verbose=false) const
Halide::Internal::Autoscheduler::State::contains_store_at
bool contains_store_at(const set< const FunctionDAG::Node * > &outermost_store_at, const IntrusivePtr< const LoopNest > &parent) const
Halide::Internal::Autoscheduler::State::mark_gpu_blocks
void mark_gpu_blocks(LoopNest::StageScheduleState *state, Stage &stage, const vector< VarOrRVar > &parallel_vars, const vector< int64_t > &parallel_extents) const
Halide::Internal::Autoscheduler::State::get_root_for_features
IntrusivePtr< const LoopNest > get_root_for_features(const Anderson2021Params &params, const Target &target) const
Halide::Internal::Autoscheduler::State::calculate_cost
bool calculate_cost(const FunctionDAG &dag, const Anderson2021Params &params, const Target &target, CostModel *cost_model, Statistics &stats, bool verbose=false)
Halide::Internal::Autoscheduler::State::create_feature_root
LoopNest * create_feature_root(const PostCreateMutator &post_create_mutator) const
Definition: State.h:106
Halide::Internal::Autoscheduler::State::deepest_common_ancestor
const LoopNest * deepest_common_ancestor(const LoopNestMap &parent, const LoopNest *a, const LoopNest *b) const
Halide::Internal::Autoscheduler::State::State
State()=default
Halide::Internal::Autoscheduler::State::make_child
IntrusivePtr< State > make_child() const
Halide::Internal::Autoscheduler::State::exceeds_local_memory_limit
bool exceeds_local_memory_limit(const Anderson2021Params &params, const Target &target) const
Halide::Internal::Autoscheduler::State::ref_count
RefCount ref_count
Definition: State.h:22
Halide::Internal::Autoscheduler::State::State
State(State &&)=delete
Halide::Internal::Autoscheduler::State::add_to_always_consider_inline_options
void add_to_always_consider_inline_options(const FunctionDAG::Node *node)
Halide::Internal::Autoscheduler::State::apply_schedule
void apply_schedule(const FunctionDAG &dag, const Anderson2021Params &params, const Target &target)
Halide::Internal::Autoscheduler::State::has_loop_nest_without_thread_loops
bool has_loop_nest_without_thread_loops() const
Halide::Internal::Autoscheduler::State::cost_per_stage
std::vector< double > cost_per_stage
Definition: State.h:80
Halide::Internal::Autoscheduler::Statistics
Definition: Statistics.h:65
Halide::Internal::IntrusivePtr
Intrusive shared pointers have a reference count (a RefCount object) stored in the class itself.
Definition: IntrusivePtr.h:71
Halide::Target
A struct representing a target machine and os to generate code for.
Definition: Target.h:19