Loading...
Searching...
No Matches
LoopNest.h
Go to the documentation of this file.
/** This file defines the LoopNest, which is our
* representation of a Halide schedule, and contains methods to
* generate candidates for scheduling as well as extract a
* featurization that can be used to cost each candidate. */
#ifndef LOOP_NEST_H
#define LOOP_NEST_H
#include "FunctionDAG.h"
#include "PerfectHashMap.h"
#include <map>
#include <set>
#include <utility>
#include <vector>
namespace Halide {
namespace Internal {
namespace Autoscheduler {
template<typename T>
using NodeMap = PerfectHashMap<FunctionDAG::Node, T>;
template<typename T>
using StageMap = PerfectHashMap<FunctionDAG::Node::Stage, T>;
// Given a multi-dimensional box of dimensionality d, generate a list
// of candidate tile sizes for it, logarithmically spacing the sizes
// using the given factor. If 'allow_splits' is false, every dimension
// must either be one, or the full extent of the box. This function is
// used to generate candidate tilings when tiling for
// producer-consumer fusion, or tiling for parallelism.
std::vector<std::vector<int64_t>> generate_tilings(const vector<int64_t> &s, int d, int factor, bool allow_splits);
struct LoopNest {
mutable RefCount ref_count;
// The extents of this loop. Put another way, the number of tiles,
// not the size of each tile.
std::vector<int64_t> size;
// The nodes inside the loop body
std::vector<IntrusivePtr<const LoopNest>> children;
// Funcs inlined into this inner loop, and the number of times
// each is called. Only valid if children is empty.
NodeMap<int64_t> inlined;
// Funcs stored inside this loop
std::set<const FunctionDAG::Node *> store_at;
// The total bounds required of any given Func over all iterations
// of this loop. In the paper, this is represented using the
// little boxes to the left of the loop nest tree figures.
mutable NodeMap<Bound> bounds;
// The Func this loop nest belongs to
const FunctionDAG::Node *node = nullptr;
// The stage of the Func
const FunctionDAG::Node::Stage *stage = nullptr;
// Is this the innermost loop of this func (the SIMD loop)?
bool innermost = false;
// Are we permitted to tile this loop?
bool tileable = false;
// Is this the parallel outer loop?
bool parallel = false;
// What dimension is this Func vectorized over, in terms of the pure args of the Func?
int vector_dim = -1;
// Which loop corresponds to the innermost storage dimension and will be vectorized. -1 means none of them.
int vectorized_loop_index = -1;
void copy_from(const LoopNest &n);
static void hash_combine(uint64_t &h, uint64_t next) {
// From boost
h ^= (next + 0x9e3779b9 + (h << 6) + (h >> 2));
}
// Hash the loop structure and sizes up to a fixed depth. This is
// used as the hash function for the coarse-to-fine beam search in
// the paper.
void structural_hash(uint64_t &h, int depth) const;
// How many funcs are scheduled inside this loop level. Used in
// the structural hash.
size_t funcs_realized_or_inlined() const {
size_t count = inlined.size() + store_at.size();
for (const auto &c : children) {
count += c->funcs_realized_or_inlined();
}
return count;
}
// All of a stage's interesting locations in the loop nest. Used to help compute the featurization of a stage.
struct Sites {
const LoopNest *compute = nullptr; // Its containing compute_at site
const LoopNest *store = nullptr; // Its containing store_at site
const LoopNest *produce = nullptr; // Its own outermost node
const LoopNest *innermost = nullptr; // Its innermost node - usually a SIMD loop
const LoopNest *task = nullptr; // The parallel for loop it belongs to
bool inlined = false; // Is the Func inlined?
// Used for caching features/feature intermediates.
uint64_t hash_of_producers_stored_at_root;
};
// Compute all the sites of interest for each pipeline stage
void get_sites(StageMap<Sites> &sites,
const LoopNest *task = nullptr,
const LoopNest *parent = nullptr) const;
// A helper for the working_set_at_task feature. Most features are
// computed in the recursive pass 'compute_features' below, but
// this one must be done in a second separate recursive pass.
void set_working_set_at_task_feature(int64_t working_set,
StageMap<ScheduleFeatures> *features) const {
for (const auto &c : children) {
c->set_working_set_at_task_feature(working_set, features);
features->get(c->stage).working_set_at_task = working_set;
}
}
// Do a recursive walk over the loop nest computing features to feed the cost model.
void compute_features(const FunctionDAG &dag,
const Adams2019Params &params,
const StageMap<Sites> &sites,
int64_t instances,
int64_t parallelism,
const LoopNest *parent,
const LoopNest *grandparent,
const LoopNest &root,
int64_t *working_set,
StageMap<ScheduleFeatures> *features,
bool use_cached_features) const;
bool is_root() const {
// The root is the sole node without a Func associated with
// it.
return node == nullptr;
}
// Set the region required of a Func at this site.
const Bound &set_bounds(const FunctionDAG::Node *f, BoundContents *b) const {
return bounds.emplace(f, b);
}
// Get the region required of a Func at this site, from which we
// know what region would be computed if it were scheduled here,
// and what its loop nest would be.
const Bound &get_bounds(const FunctionDAG::Node *f) const;
// Recursively print a loop nest representation to stderr
void dump(std::ostream &os, string prefix, const LoopNest *parent) const;
// Does this loop nest access the given Func
bool calls(const FunctionDAG::Node *f) const;
// What is the maximum number of inlined calls to a Func that
// occur within this loop. Used to prune states that would
// generate too much code.
int64_t max_inlined_calls() const;
// Does this loop nest access an input buffer? Used to select
// trail strategies when splitting loops. We don't want to read
// out of bounds on inputs, even if we don't intend to use the
// values read. It could create annoying assertion failures for
// the user. It's OK to read out of range of the values computed
// on internal Funcs though. Allocation bounds inference just pads
// out the bounds so that it won't fault.
bool accesses_input_buffer() const;
// Does this loop nest contain a computation of the given Func.
bool computes(const FunctionDAG::Node *f) const;
// Above here most methods query the loop nest. Below we have
// methods that mutate the loop nest.
// Inline a Func into all consumers within this loop.
void inline_func(const FunctionDAG::Node *f);
// Compute a Func at this site.
void compute_here(const FunctionDAG::Node *f, bool tileable, int v, const Adams2019Params &params);
// Parallelize this loop according to the given tiling.
IntrusivePtr<const LoopNest> parallelize_in_tiles(const Adams2019Params &params,
const vector<int64_t> &tiling,
const LoopNest *parent) const;
// Return all possible ways to compute f in tiles somewhere within
// this loop nest.
std::vector<IntrusivePtr<const LoopNest>> compute_in_tiles(const FunctionDAG::Node *f,
const LoopNest *parent,
const Adams2019Params &params,
int v,
bool in_realization) const;
// Below here we have methods that apply a schedule to a Halide pipeline.
// A model of the state of the loop nest of a Func while applying
// Halide's scheduling directives.
// Note that StageScheduleState is movable-but-not-copyable thanks
// to its ostringstream member.
struct StageScheduleState {
// How much parallelism do we need to exploit with this Func?
double num_cores = 0;
// Which storage dimension is vectorized? We need to reorder it innermost
int vector_dim = -1;
int vectorized_loop_index = -1;
// The various Vars and RVars used for scheduling a Func.
struct FuncVar {
// The top-level var or rvar this was split off from
VarOrRVar orig;
// This var.
VarOrRVar var;
// Source code to access this Var/RVar. Used for printing
// valid Halide source for this schedule.
string accessor;
// Our estimate of the extent of this var. This is exact
// when constant_extent flag is true.
int64_t extent = 0;
// Which index in the symbolic loop nest does this var
// belong to.
size_t index = 0;
// Some flags.
bool innermost_pure_dim = false,
outermost = false,
parallel = false,
exists = false,
pure = false,
constant_extent = false;
FuncVar()
: orig(Var()), var(Var()) {
}
};
// In order from innermost to outermost. Each group of d is one tiling level.
std::vector<FuncVar> vars;
std::ostringstream schedule_source;
};
// Apply the schedule represented by this loop nest to a Halide pipeline.
void apply(LoopLevel here,
StageMap<std::unique_ptr<StageScheduleState>> &state_map,
double num_cores,
int depth,
const LoopNest *parent,
const LoopNest *compute_site) const;
// The below are two feature caches.
// hash of producers -> StageMap
mutable std::map<uint64_t, StageMap<StageMap<FeatureIntermediates>>> feature_intermediates_cache;
// hash of producers -> StageMap
mutable std::map<uint64_t, StageMap<ScheduleFeatures>> features_cache;
// Same as copy_from (above) but also copies the two caches.
void copy_from_including_features(const LoopNest &n);
// Loops through inlined funcs and caches the pcm found in features, into memoized_features.
void memoize_points_computed_minimum(StageMap<ScheduleFeatures> &memoized_features,
const StageMap<ScheduleFeatures> *features) const;
// Merges features_to_insert into memoized_features if it does not already exist there.
void memoize_features(StageMap<ScheduleFeatures> &memoized_features,
const StageMap<ScheduleFeatures> *features_to_insert) const;
// Recalculates working_set from cached features
void compute_working_set_from_features(int64_t *working_set,
const StageMap<ScheduleFeatures> *features) const;
// Features need to be recomputed for inlined Funcs
void recompute_inlined_features(const StageMap<Sites> &sites,
StageMap<ScheduleFeatures> *features) const;
// Create a (hopefully) unique hash of the producers.
uint64_t compute_hash_of_producers_stored_at_root(const StageMap<Sites> &sites) const;
// Gather all stages that are producers for any Func in this LoopNest.
std::vector<std::pair<int, int>> collect_producers(const StageMap<Sites> &sites) const;
// Collect all stages referenced in this LoopNest.
void collect_stages(std::set<const FunctionDAG::Node::Stage *> &stages) const;
};
// Find the deepest common ancestor of `a` and `b`.
// `parents` is a map from loop nest to (parent, depth) tuples.
// Assumes that `a` and `b` are found in `parents`, otherwise errors.
const LoopNest *deepest_common_ancestor(const std::map<const LoopNest *, std::pair<const LoopNest *, int>> &parents,
const LoopNest *a, const LoopNest *b);
// Compute the parent and depth of every loop nest node.
// Stores in `parents` the children of `here` (keys) to tuples of (here, depth).
// Recurses on all children of `here`.
void compute_loop_nest_parents(std::map<const LoopNest *, std::pair<const LoopNest *, int>> &parents,
const LoopNest *here, int depth);
} // namespace Autoscheduler
} // namespace Internal
} // namespace Halide
#endif // LOOP_NEST_H
PerfectHashMap.h
FunctionDAG.h
Halide::Internal::RefCount
A class representing a reference count to be used with IntrusivePtr.
Definition IntrusivePtr.h:19
Halide::LoopLevel
A reference to a site in a Halide statement at the top of the body of a particular for loop.
Definition Schedule.h:203
Halide::Var
A Halide variable, to be used when defining functions.
Definition Var.h:19
PerfectHashMap
Definition PerfectHashMap.h:38
Halide::Internal::Autoscheduler
Definition AutoSchedule.h:10
Halide::Internal::Autoscheduler::StageMap
PerfectHashMap< FunctionDAG::Node::Stage, T > StageMap
Definition LoopNest.h:24
Halide::Internal::Autoscheduler::Bound
IntrusivePtr< const BoundContents > Bound
Definition FunctionDAG.h:362
Halide::Internal::Autoscheduler::generate_tilings
std::vector< std::vector< int64_t > > generate_tilings(const vector< int64_t > &s, int d, int factor, bool allow_splits)
Halide::Internal::Autoscheduler::deepest_common_ancestor
const LoopNest * deepest_common_ancestor(const std::map< const LoopNest *, std::pair< const LoopNest *, int > > &parents, const LoopNest *a, const LoopNest *b)
Halide::Internal::Autoscheduler::NodeMap
PerfectHashMap< FunctionDAG::Node, T > NodeMap
Definition LoopNest.h:21
Halide::Internal::Autoscheduler::compute_loop_nest_parents
void compute_loop_nest_parents(std::map< const LoopNest *, std::pair< const LoopNest *, int > > &parents, const LoopNest *here, int depth)
Halide::Internal
Definition AbstractGenerator.h:24
Halide
This file defines the class FunctionDAG, which is our representation of a Halide pipeline,...
Definition AbstractGenerator.h:19
uint64_t
unsigned __INT64_TYPE__ uint64_t
Definition runtime_internal.h:23
int64_t
signed __INT64_TYPE__ int64_t
Definition runtime_internal.h:22
Halide::Internal::Autoscheduler::Adams2019Params
Definition CostModel.h:19
Halide::Internal::Autoscheduler::BoundContents
Definition FunctionDAG.h:274
Halide::Internal::Autoscheduler::FunctionDAG::Node::Stage
Definition FunctionDAG.h:448
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::Sites::hash_of_producers_stored_at_root
uint64_t hash_of_producers_stored_at_root
Definition LoopNest.h:109
Halide::Internal::Autoscheduler::LoopNest::Sites::task
const LoopNest * task
Definition LoopNest.h:105
Halide::Internal::Autoscheduler::LoopNest::Sites::innermost
const LoopNest * innermost
Definition LoopNest.h:104
Halide::Internal::Autoscheduler::LoopNest::Sites::store
const LoopNest * store
Definition LoopNest.h:102
Halide::Internal::Autoscheduler::LoopNest::Sites::produce
const LoopNest * produce
Definition LoopNest.h:103
Halide::Internal::Autoscheduler::LoopNest::Sites::inlined
bool inlined
Definition LoopNest.h:106
Halide::Internal::Autoscheduler::LoopNest::Sites::compute
const LoopNest * compute
Definition LoopNest.h:101
Halide::Internal::Autoscheduler::LoopNest::StageScheduleState::FuncVar::FuncVar
FuncVar()
Definition LoopNest.h:244
Halide::Internal::Autoscheduler::LoopNest::StageScheduleState::FuncVar::innermost_pure_dim
bool innermost_pure_dim
Definition LoopNest.h:238
Halide::Internal::Autoscheduler::LoopNest::StageScheduleState::FuncVar::index
size_t index
Definition LoopNest.h:235
Halide::Internal::Autoscheduler::LoopNest::StageScheduleState::FuncVar::pure
bool pure
Definition LoopNest.h:242
Halide::Internal::Autoscheduler::LoopNest::StageScheduleState::FuncVar::orig
VarOrRVar orig
Definition LoopNest.h:220
Halide::Internal::Autoscheduler::LoopNest::StageScheduleState::FuncVar::constant_extent
bool constant_extent
Definition LoopNest.h:243
Halide::Internal::Autoscheduler::LoopNest::StageScheduleState::FuncVar::parallel
bool parallel
Definition LoopNest.h:240
Halide::Internal::Autoscheduler::LoopNest::StageScheduleState::FuncVar::exists
bool exists
Definition LoopNest.h:241
Halide::Internal::Autoscheduler::LoopNest::StageScheduleState::FuncVar::outermost
bool outermost
Definition LoopNest.h:239
Halide::Internal::Autoscheduler::LoopNest::StageScheduleState::FuncVar::var
VarOrRVar var
Definition LoopNest.h:223
Halide::Internal::Autoscheduler::LoopNest::StageScheduleState::FuncVar::extent
int64_t extent
Definition LoopNest.h:231
Halide::Internal::Autoscheduler::LoopNest::StageScheduleState::FuncVar::accessor
string accessor
Definition LoopNest.h:227
Halide::Internal::Autoscheduler::LoopNest::StageScheduleState
Definition LoopNest.h:209
Halide::Internal::Autoscheduler::LoopNest::StageScheduleState::num_cores
double num_cores
Definition LoopNest.h:211
Halide::Internal::Autoscheduler::LoopNest::StageScheduleState::vars
std::vector< FuncVar > vars
Definition LoopNest.h:250
Halide::Internal::Autoscheduler::LoopNest::StageScheduleState::vector_dim
int vector_dim
Definition LoopNest.h:214
Halide::Internal::Autoscheduler::LoopNest::StageScheduleState::vectorized_loop_index
int vectorized_loop_index
Definition LoopNest.h:215
Halide::Internal::Autoscheduler::LoopNest::StageScheduleState::schedule_source
std::ostringstream schedule_source
Definition LoopNest.h:252
Halide::Internal::Autoscheduler::LoopNest
Definition LoopNest.h:34
Halide::Internal::Autoscheduler::LoopNest::parallelize_in_tiles
IntrusivePtr< const LoopNest > parallelize_in_tiles(const Adams2019Params &params, const vector< int64_t > &tiling, const LoopNest *parent) const
Halide::Internal::Autoscheduler::LoopNest::vectorized_loop_index
int vectorized_loop_index
Definition LoopNest.h:75
Halide::Internal::Autoscheduler::LoopNest::is_root
bool is_root() const
Definition LoopNest.h:141
Halide::Internal::Autoscheduler::LoopNest::features
std::map< uint64_t, StageMap< ScheduleFeatures > > features
Definition LoopNest.h:146
Halide::Internal::Autoscheduler::LoopNest::node
const FunctionDAG::Node * node
Definition LoopNest.h:57
Halide::Internal::Autoscheduler::LoopNest::recompute_inlined_features
void recompute_inlined_features(const StageMap< Sites > &sites, StageMap< ScheduleFeatures > *features) const
Halide::Internal::Autoscheduler::LoopNest::inline_func
void inline_func(const FunctionDAG::Node *f)
Halide::Internal::Autoscheduler::LoopNest::compute_in_tiles
std::vector< IntrusivePtr< const LoopNest > > compute_in_tiles(const FunctionDAG::Node *f, const LoopNest *parent, const Adams2019Params &params, int v, bool in_realization) const
Halide::Internal::Autoscheduler::LoopNest::collect_stages
void collect_stages(std::set< const FunctionDAG::Node::Stage * > &stages) const
Halide::Internal::Autoscheduler::LoopNest::accesses_input_buffer
bool accesses_input_buffer() const
Halide::Internal::Autoscheduler::LoopNest::size
std::vector< int64_t > size
Definition LoopNest.h:39
Halide::Internal::Autoscheduler::LoopNest::get_sites
void get_sites(StageMap< Sites > &sites, const LoopNest *task=nullptr, const LoopNest *parent=nullptr) const
Halide::Internal::Autoscheduler::LoopNest::get_bounds
const Bound & get_bounds(const FunctionDAG::Node *f) const
Halide::Internal::Autoscheduler::LoopNest::compute_features
void compute_features(const FunctionDAG &dag, const Adams2019Params &params, const StageMap< Sites > &sites, int64_t instances, int64_t parallelism, const LoopNest *parent, const LoopNest *grandparent, const LoopNest &root, int64_t *working_set, StageMap< ScheduleFeatures > *features, bool use_cached_features) const
Halide::Internal::Autoscheduler::LoopNest::memoize_features
void memoize_features(StageMap< ScheduleFeatures > &memoized_features, const StageMap< ScheduleFeatures > *features_to_insert) const
Halide::Internal::Autoscheduler::LoopNest::vector_dim
int vector_dim
Definition LoopNest.h:72
Halide::Internal::Autoscheduler::LoopNest::features_cache
std::map< uint64_t, StageMap< ScheduleFeatures > > features_cache
Definition LoopNest.h:267
Halide::Internal::Autoscheduler::LoopNest::computes
bool computes(const FunctionDAG::Node *f) const
Halide::Internal::Autoscheduler::LoopNest::tileable
bool tileable
Definition LoopNest.h:66
Halide::Internal::Autoscheduler::LoopNest::children
std::vector< IntrusivePtr< const LoopNest > > children
Definition LoopNest.h:42
Halide::Internal::Autoscheduler::LoopNest::set_working_set_at_task_feature
void set_working_set_at_task_feature(int64_t working_set, StageMap< ScheduleFeatures > *features) const
Definition LoopNest.h:120
Halide::Internal::Autoscheduler::LoopNest::copy_from
void copy_from(const LoopNest &n)
Halide::Internal::Autoscheduler::LoopNest::collect_producers
std::vector< std::pair< int, int > > collect_producers(const StageMap< Sites > &sites) const
Halide::Internal::Autoscheduler::LoopNest::set_bounds
const Bound & set_bounds(const FunctionDAG::Node *f, BoundContents *b) const
Definition LoopNest.h:148
Halide::Internal::Autoscheduler::LoopNest::memoize_points_computed_minimum
void memoize_points_computed_minimum(StageMap< ScheduleFeatures > &memoized_features, const StageMap< ScheduleFeatures > *features) const
Halide::Internal::Autoscheduler::LoopNest::hash_combine
static void hash_combine(uint64_t &h, uint64_t next)
Definition LoopNest.h:79
Halide::Internal::Autoscheduler::LoopNest::dump
void dump(std::ostream &os, string prefix, const LoopNest *parent) const
Halide::Internal::Autoscheduler::LoopNest::ref_count
RefCount ref_count
Definition LoopNest.h:35
Halide::Internal::Autoscheduler::LoopNest::bounds
NodeMap< Bound > bounds
Definition LoopNest.h:54
Halide::Internal::Autoscheduler::LoopNest::compute_hash_of_producers_stored_at_root
uint64_t compute_hash_of_producers_stored_at_root(const StageMap< Sites > &sites) const
Halide::Internal::Autoscheduler::LoopNest::max_inlined_calls
int64_t max_inlined_calls() const
Halide::Internal::Autoscheduler::LoopNest::stage
const FunctionDAG::Node::Stage * stage
Definition LoopNest.h:60
Halide::Internal::Autoscheduler::LoopNest::innermost
bool innermost
Definition LoopNest.h:63
Halide::Internal::Autoscheduler::LoopNest::compute_here
void compute_here(const FunctionDAG::Node *f, bool tileable, int v, const Adams2019Params &params)
Halide::Internal::Autoscheduler::LoopNest::parallel
bool parallel
Definition LoopNest.h:69
Halide::Internal::Autoscheduler::LoopNest::calls
bool calls(const FunctionDAG::Node *f) const
Halide::Internal::Autoscheduler::LoopNest::copy_from_including_features
void copy_from_including_features(const LoopNest &n)
Halide::Internal::Autoscheduler::LoopNest::structural_hash
void structural_hash(uint64_t &h, int depth) const
Halide::Internal::Autoscheduler::LoopNest::funcs_realized_or_inlined
size_t funcs_realized_or_inlined() const
Definition LoopNest.h:91
Halide::Internal::Autoscheduler::LoopNest::store_at
std::set< const FunctionDAG::Node * > store_at
Definition LoopNest.h:49
Halide::Internal::Autoscheduler::LoopNest::feature_intermediates_cache
std::map< uint64_t, StageMap< StageMap< FeatureIntermediates > > > feature_intermediates_cache
Definition LoopNest.h:265
Halide::Internal::Autoscheduler::LoopNest::apply
void apply(LoopLevel here, StageMap< std::unique_ptr< StageScheduleState > > &state_map, double num_cores, int depth, const LoopNest *parent, const LoopNest *compute_site) const
Halide::Internal::Autoscheduler::LoopNest::inlined
NodeMap< int64_t > inlined
Definition LoopNest.h:46
Halide::Internal::Autoscheduler::LoopNest::compute_working_set_from_features
void compute_working_set_from_features(int64_t *working_set, const StageMap< ScheduleFeatures > *features) const
Halide::Internal::IntrusivePtr
Intrusive shared pointers have a reference count (a RefCount object) stored in the class itself.
Definition IntrusivePtr.h:71
Halide::VarOrRVar
A class that can represent Vars or RVars.
Definition Func.h:29