Halide 14.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 "Cache.h"
6#include "CostModel.h"
7#include "DefaultCostModel.h"
8#include "Featurization.h"
9#include "Halide.h"
10#include "LoopNest.h"
11#include "PerfectHashMap.h"
12#include <map>
13#include <utility>
14
15namespace Halide {
16namespace Internal {
17namespace Autoscheduler {
18
19// A struct representing an intermediate state in the tree search.
20// It represents a partial schedule for some pipeline.
21struct State {
23 // The LoopNest this state corresponds to.
25 // The parent that generated this state.
27 // Cost of this state, as evaluated by the cost model.
28 double cost = 0;
29 // Number of decisions made at this state (used for finding which DAG node to schedule).
31 // Penalization is determined based on structural hash during beam search.
32 bool penalized = false;
33
34 // The C++ source code of the generated schedule for this State.
35 // Computed if `apply_schedule` is called.
37
38 // The number of times a cost is enqueued into the cost model,
39 // for all states.
41
42 State() = default;
43 State(const State &) = delete;
44 State(State &&) = delete;
45 void operator=(const State &) = delete;
46 void operator=(State &&) = delete;
47
48 // Compute a structural hash based on depth and num_decisions_made.
49 // Defers to root->structural_hash().
50 uint64_t structural_hash(int depth) const;
51
52 // Compute the featurization of this state (based on `root`),
53 // and store features in `features`. Defers to `root->compute_features()`.
55 const MachineParams &params,
57 const CachingOptions &cache_options);
58
59 // Calls `compute_featurization` and prints those features to `out`.
61 const MachineParams &params,
62 const CachingOptions &cache_options,
63 std::ostream &out);
64
65 // Performs some pruning to decide if this state is worth queuing in
66 // the cost_model. If it is, calls `cost_model->enqueue` and returns true,
67 // otherwise sets `cost` equal to a large value and returns false.
68 bool calculate_cost(const FunctionDAG &dag, const MachineParams &params,
69 CostModel *cost_model, const CachingOptions &cache_options,
70 int64_t memory_limit, bool verbose = false);
71
72 // Make a child copy of this state. The loop nest is const (we
73 // make mutated copies of it, rather than mutating it), so we can
74 // continue to point to the same one and so this is a cheap
75 // operation.
77
78 // Generate the successor states to this state.
79 // If they are not pruned by `calculate_cost()`,
80 // then calls `accept_child()` on them.
82 const MachineParams &params,
83 CostModel *cost_model,
84 int64_t memory_limit,
85 std::function<void(IntrusivePtr<State> &&)> &accept_child,
86 Cache *cache) const;
87
88 // Dumps cost, the `root` LoopNest, and then `schedule_source` to `aslog(0)`.
89 void dump() const;
90
91 // Apply the schedule represented by this state to a Halide
92 // Pipeline. Also generate source code for the schedule for the
93 // user to copy-paste to freeze this schedule as permanent artifact.
94 // Also fills `schedule_source`.
95 void apply_schedule(const FunctionDAG &dag, const MachineParams &params);
96};
97
98} // namespace Autoscheduler
99} // namespace Internal
100} // namespace Halide
101
102#endif // STATE_H
A class representing a reference count to be used with IntrusivePtr.
Definition: IntrusivePtr.h:19
This file defines the class FunctionDAG, which is our representation of a Halide pipeline,...
@ Internal
Not visible externally, similar to 'static' linkage in C.
unsigned __INT64_TYPE__ uint64_t
signed __INT64_TYPE__ int64_t
void apply_schedule(const FunctionDAG &dag, const MachineParams &params)
void operator=(const State &)=delete
void save_featurization(const FunctionDAG &dag, const MachineParams &params, const CachingOptions &cache_options, std::ostream &out)
uint64_t structural_hash(int depth) const
bool calculate_cost(const FunctionDAG &dag, const MachineParams &params, CostModel *cost_model, const CachingOptions &cache_options, int64_t memory_limit, bool verbose=false)
void generate_children(const FunctionDAG &dag, const MachineParams &params, CostModel *cost_model, int64_t memory_limit, std::function< void(IntrusivePtr< State > &&)> &accept_child, Cache *cache) const
IntrusivePtr< const LoopNest > root
Definition: State.h:24
IntrusivePtr< State > make_child() const
void compute_featurization(const FunctionDAG &dag, const MachineParams &params, StageMap< ScheduleFeatures > *features, const CachingOptions &cache_options)
IntrusivePtr< const State > parent
Definition: State.h:26
Intrusive shared pointers have a reference count (a RefCount object) stored in the class itself.
Definition: IntrusivePtr.h:68
A struct representing the machine parameters to generate the auto-scheduled code for.
Definition: Pipeline.h:33