Loading...
Searching...
No Matches
LoopNestParser.h
Go to the documentation of this file.
#ifndef LOOP_NEST_PARSER_H
#define LOOP_NEST_PARSER_H
#include <fstream>
#include <sstream>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include "ASLog.h"
#include "FunctionDAG.h"
namespace Halide {
namespace Internal {
namespace Autoscheduler {
class LoopNestParser {
void parse(const std::vector<std::string> &loop_nest) {
std::unordered_map<std::string, std::vector<std::string>> stage_to_loop_nest;
for (const auto &line : loop_nest) {
if (line.empty()) {
continue;
}
if (line.at(0) == '#') {
continue;
}
std::istringstream iss(line);
std::vector<std::string> tokens{
std::istream_iterator<std::string>(iss),
std::istream_iterator<std::string>()};
std::string stage = tokens.at(0);
bool is_inlined = tokens.at(0) == "inlined:";
if (tokens.at(0) == "realize:" || is_inlined) {
stage = tokens.at(1);
}
if (stage == "gpu_none") {
continue;
}
all_stages.insert(stage);
if (is_inlined) {
inlined.insert(stage);
continue;
}
if (tokens.back() == "gpu_none") {
partially_scheduled.insert(stage);
}
if (line.at(0) != ' ' && compute_root_stages.count(stage) == 0) {
compute_root_stages[stage] = -1;
}
if (tokens.back() == "gpu_simd" && compute_root_stages.count(stage) == 1 && compute_root_stages[stage] == -1) {
const std::string &vector_dim = tokens[tokens.size() - 3];
compute_root_stages[stage] = std::stoi(vector_dim.substr(0, vector_dim.size() - 1));
}
if (partially_scheduled.count(stage) == 0) {
stage_to_loop_nest[stage].push_back(line);
}
}
for (const auto &entry : stage_to_loop_nest) {
std::string loop_nest = "";
for (const auto &line : entry.second) {
loop_nest += line + "\n";
}
per_stage_loop_nests[entry.first] = loop_nest;
}
// If a stage appears in a 'realize: ' line but nowhere else, remove it
std::vector<std::string> to_remove;
for (const auto &entry : compute_root_stages) {
if (entry.second == -1) {
to_remove.push_back(entry.first);
}
}
for (const auto &key : to_remove) {
compute_root_stages.erase(key);
partially_scheduled.erase(key);
all_stages.erase(key);
per_stage_loop_nests.erase(key);
}
}
std::vector<std::string> loop_nest;
std::unordered_map<std::string, std::string> per_stage_loop_nests;
std::unordered_set<std::string> inlined;
std::unordered_set<std::string> partially_scheduled;
std::unordered_map<std::string, int> compute_root_stages;
std::unordered_set<std::string> all_stages;
public:
explicit LoopNestParser(const std::vector<std::string> &loop_nest)
: loop_nest{loop_nest} {
parse(loop_nest);
}
void dump() const {
aslog(1) << "All stages:\n";
for (const auto &s : all_stages) {
aslog(1) << s << "\n";
}
aslog(1) << "\ncompute_root stages:\n";
for (const auto &s : compute_root_stages) {
aslog(1) << s.first << " with vector_dim = " << s.second << "\n";
}
aslog(1) << "\nPartially scheduled stages:\n";
for (const auto &s : partially_scheduled) {
aslog(1) << s << " with vector_dim = " << compute_root_stages.at(s) << "\n";
}
aslog(1) << "\nInlined stages:\n";
for (const auto &s : inlined) {
aslog(1) << s << "\n";
}
aslog(1) << "\nFull loop nest:\n";
for (const auto &s : loop_nest) {
aslog(1) << s << "\n";
}
aslog(1) << "\n";
}
bool is_in_partial_schedule(const FunctionDAG::Node *node) const {
return node && all_stages.count(node->func.name()) > 0;
}
bool contains_sub_loop_nest_for_shared_stages(const LoopNestParser &other) const {
return contains_sub_loop_nest(other, true);
}
// 'only_consider_shared_stages': check if 'other' is contained in this loop
// nest, but ignore stages that are present in 'other' but not present in
// this loop nest
bool contains_sub_loop_nest(const LoopNestParser &other, bool only_consider_shared_stages = false) const {
for (const auto &stage : other.all_stages) {
if (all_stages.count(stage) == 0) {
if (only_consider_shared_stages) {
continue;
}
return false;
}
if (other.partially_scheduled.count(stage) == 1) {
if (compute_root_stages.count(stage) == 0) {
return false;
}
return other.compute_root_stages.at(stage) == compute_root_stages.at(stage);
}
if (other.inlined.count(stage) > 0) {
if (inlined.count(stage) == 0) {
return false;
}
continue;
} else if (inlined.count(stage) > 0) {
return false;
}
if (other.per_stage_loop_nests.at(stage) != per_stage_loop_nests.at(stage)) {
return false;
}
}
return true;
}
static LoopNestParser from_string(const std::string &str) {
std::istringstream in(str);
std::string line;
std::vector<std::string> loop_nest;
while (std::getline(in, line)) {
loop_nest.push_back(line);
}
return LoopNestParser(loop_nest);
}
static std::unique_ptr<LoopNestParser> from_file(const std::string &filename) {
std::ifstream file(filename);
std::string line;
std::vector<std::string> loop_nest;
while (std::getline(file, line)) {
loop_nest.push_back(line);
}
return std::make_unique<LoopNestParser>(loop_nest);
}
};
} // namespace Autoscheduler
} // namespace Internal
} // namespace Halide
#endif
ASLog.h
FunctionDAG.h
Halide::Internal::Autoscheduler::LoopNestParser::dump
void dump() const
Definition LoopNestParser.h:108
Halide::Internal::Autoscheduler::LoopNestParser::LoopNestParser
LoopNestParser(const std::vector< std::string > &loop_nest)
Definition LoopNestParser.h:103
Halide::Internal::Autoscheduler::LoopNestParser::contains_sub_loop_nest_for_shared_stages
bool contains_sub_loop_nest_for_shared_stages(const LoopNestParser &other) const
Definition LoopNestParser.h:140
Halide::Internal::Autoscheduler::LoopNestParser::from_string
static LoopNestParser from_string(const std::string &str)
Definition LoopNestParser.h:181
Halide::Internal::Autoscheduler::LoopNestParser::contains_sub_loop_nest
bool contains_sub_loop_nest(const LoopNestParser &other, bool only_consider_shared_stages=false) const
Definition LoopNestParser.h:147
Halide::Internal::Autoscheduler::LoopNestParser::from_file
static std::unique_ptr< LoopNestParser > from_file(const std::string &filename)
Definition LoopNestParser.h:193
Halide::Internal::Autoscheduler::LoopNestParser::is_in_partial_schedule
bool is_in_partial_schedule(const FunctionDAG::Node *node) const
Definition LoopNestParser.h:136
Halide::Internal::Function::name
const std::string & name() const
Get the name of the function.
Halide::Internal::aslog
Definition ASLog.h:16
Halide::Internal::Autoscheduler
Definition AutoSchedule.h:10
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
Halide::Internal::Autoscheduler::FunctionDAG::Node
Definition FunctionDAG.h:378
Halide::Internal::Autoscheduler::FunctionDAG::Node::func
Function func
Definition FunctionDAG.h:383