From 36246588aa887eb83c978159890f850cf38ea12f Mon Sep 17 00:00:00 2001 From: Superjomn Date: Tue, 4 Feb 2020 20:05:03 +0800 Subject: [PATCH] init schedule --- cinn/poly/CMakeLists.txt | 2 + cinn/poly/ast_gen.cc | 1 + cinn/poly/ast_gen.h | 1 + cinn/poly/element.cc | 2 + cinn/poly/element.h | 8 +++ cinn/poly/isl_utils.cc | 8 +++ cinn/poly/isl_utils.h | 3 + cinn/poly/schedule.cc | 63 ++++++++++++++++++++ cinn/poly/schedule.h | 120 +++++++++++++++++++++++++++++++++++++++ 9 files changed, 208 insertions(+) create mode 100644 cinn/poly/ast_gen.cc create mode 100644 cinn/poly/ast_gen.h create mode 100644 cinn/poly/schedule.cc create mode 100644 cinn/poly/schedule.h diff --git a/cinn/poly/CMakeLists.txt b/cinn/poly/CMakeLists.txt index 3621e71d010c0..cc2261c0f5159 100644 --- a/cinn/poly/CMakeLists.txt +++ b/cinn/poly/CMakeLists.txt @@ -4,6 +4,8 @@ cc_library(poly SRCS map.cc element.cc isl_utils.cc + schedule.cc + ast_gen.cc DEPS common) cc_test(test_poly_element SRCS element_test.cc DEPS poly) diff --git a/cinn/poly/ast_gen.cc b/cinn/poly/ast_gen.cc new file mode 100644 index 0000000000000..8b137891791fe --- /dev/null +++ b/cinn/poly/ast_gen.cc @@ -0,0 +1 @@ + diff --git a/cinn/poly/ast_gen.h b/cinn/poly/ast_gen.h new file mode 100644 index 0000000000000..6f70f09beec22 --- /dev/null +++ b/cinn/poly/ast_gen.h @@ -0,0 +1 @@ +#pragma once diff --git a/cinn/poly/element.cc b/cinn/poly/element.cc index 1a5d7eff992ff..2efa689dceaa0 100644 --- a/cinn/poly/element.cc +++ b/cinn/poly/element.cc @@ -99,5 +99,7 @@ std::string OuterName(const std::string &name) { return name + "_outer"; } std::string InnerName(const Iterator &iterator) { return InnerName(iterator.id); } std::string OuterName(const Iterator &iterator) { return OuterName(iterator.id); } +const char *Element::id() const { return isl_set_get_tuple_name(domain_.get()); } + } // namespace poly } // namespace cinn diff --git a/cinn/poly/element.h b/cinn/poly/element.h index d823580c75516..00c909dda1fdd 100644 --- a/cinn/poly/element.h +++ b/cinn/poly/element.h @@ -21,6 +21,11 @@ class Element { public: explicit Element(isl::set domain); + /** + * The id of this element, should be unique across the schedule. + */ + const char* id() const; + /** * Split the loop level of into two new loop levels. * @param level the level to split. @@ -61,6 +66,9 @@ class Element { */ Iterator Fuse(const Iterator& level0, const Iterator& level1); + const isl::set& domain() const { return domain_; } + const isl::map& schedule() const { return schedule_; } + private: /** * Initialize with an identity schedule. diff --git a/cinn/poly/isl_utils.cc b/cinn/poly/isl_utils.cc index bb36730c78d5a..a341726bb70c4 100644 --- a/cinn/poly/isl_utils.cc +++ b/cinn/poly/isl_utils.cc @@ -11,5 +11,13 @@ std::vector GetDimNames(const isl::set &x) { return res; } +std::vector poly::GetDimNames(const isl::map &x, isl_dim_type dim_type) { + std::vector res; + for (int i = 0; i < isl_map_dim(x.get(), dim_type); i++) { + res.push_back(isl_map_get_dim_name(x.get(), dim_type, i)); + } + return res; +} + } // namespace poly } // namespace cinn diff --git a/cinn/poly/isl_utils.h b/cinn/poly/isl_utils.h index 11c8db9703fff..bd4a4b4fb9807 100644 --- a/cinn/poly/isl_utils.h +++ b/cinn/poly/isl_utils.h @@ -8,7 +8,10 @@ namespace cinn { namespace poly { //! Get dimension names from isl containers. +// @{ std::vector GetDimNames(const isl::set &x); +std::vector GetDimNames(const isl::map &x, isl_dim_type dim_type); +// @} } // namespace poly } // namespace cinn diff --git a/cinn/poly/schedule.cc b/cinn/poly/schedule.cc new file mode 100644 index 0000000000000..7f07c342fc449 --- /dev/null +++ b/cinn/poly/schedule.cc @@ -0,0 +1,63 @@ +#include "cinn/poly/schedule.h" +#include "cinn/utils/string.h" + +#include + +namespace cinn { +namespace poly { + +std::string TimeSchedule::__str__() const { + CHECK(!time_dims.empty()); + + // generate range: [dup, t0, t1...] + std::vector range_dims({"dup"}); + for (int i = 0; i < time_dims.size(); i++) { + range_dims.push_back("t" + std::to_string(i)); + } + + // generate conditions + std::vector conds; + for (int i = 0; i < time_dims.size(); i++) { + conds.push_back(std::to_string(time_dims[i].time)); + conds.push_back(time_dims[i].dim); + } + + return utils::StringFormat("{ %s[%s] -> [%s]: %s", + id.c_str(), + utils::Join(domain_dims, ", ").c_str(), + utils::Join(range_dims, ", ").c_str(), + utils::Join(conds, " and ").c_str()); +} + +void Scheduler::RegisterElement(const Element &x) { + CHECK(!registration_finalized_) << "element registration has been finalized."; + space_size_ = std::max(space_size_, isl_map_dim(x.schedule().get(), isl_dim_out)); + + // Use the dimensions from element's schedule's range as the new domain dimensions because in Element, the schedule is + // like '{ S0[i,j] -> S0[i_outer, i_inner, j] }', the scheduler should schedule base on the range. + TimeSchedule schedule(GetDimNames(x.schedule(), isl_dim_out)); + schedule_.emplace(x.id(), std::move(schedule)); +} + +void Scheduler::FinalizeRegistration() { + CHECK_GT(space_size_, 0) << "No valid dimension is collected, use RegisterElement to collect some elements"; + CHECK(!schedule_.empty()) << "No valid dimension is collected, use RegisterElement to collect some elements"; + registration_finalized_ = false; + + for (auto &item : schedule_) { + item.second.ResizeTimeSpace(space_size_); + } +} + +Scheduler &Scheduler::After(const Element &a, const Element &b, int level) { + CHECK_LT(level, space_size_); + depend_flow_graph_[b.id()].depend_level[a.id()] = level; + return *this; +} + +Scheduler &Scheduler::Before(const Element &a, const Element &b, int level) { return After(b, a, level); } + +std::unordered_map Scheduler::BuildSchedule() const {} + +} // namespace poly +} // namespace cinn diff --git a/cinn/poly/schedule.h b/cinn/poly/schedule.h new file mode 100644 index 0000000000000..fbe576233c095 --- /dev/null +++ b/cinn/poly/schedule.h @@ -0,0 +1,120 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "cinn/poly/element.h" +#include "cinn/poly/isl_utils.h" +#include "cinn/poly/map.h" + +namespace cinn { +namespace poly { + +struct TimeDim { + //! time of this dimension. + int time; + //! name of this dimension. + std::string dim; + + TimeDim(std::string dim, int time) : dim(std::move(dim)), time(time) {} +}; + +struct DependFlow { + //! Map from the depended Element.id to the level. + std::unordered_map depend_level; +}; + +/** + * The range of the schedule. + */ +struct TimeSchedule { + //! ISL range format, such as '[dup, t0, t1]: dup=0 and t0=0 and t1=i]' + std::string __str__() const; + + TimeSchedule(const std::vector &dims) { + domain_dims = dims; + for (auto &dim : domain_dims) { + time_dims.emplace_back(dim, 0); + } + } + + void ResizeTimeSpace(int size) { time_dims.resize(size); } + + //! Get the isl map. + isl::map to_isl(isl::ctx ctx) const { return isl::map(ctx, __str__()); } + + std::string id; + std::vector domain_dims; + int duplicate_id{}; + std::vector time_dims; +}; + +/** + * Scheduler - Perform schedule on polyhedral model. + * It takes a normal schedule as input, and transform it into + * + */ +class Scheduler { + public: + /** + * Constructor. + * @param schedule A normal isl schedule, such as '{ S[i,j] -> [i,j] }' + * + * The schedule input can be transformed, that's ok, such as + * '{ S[i,j] -> [i_outer, i_inner, j]: i_outer=floor(i/4) and i_inner=i%4 }' + * that's OK. + */ + Scheduler() = default; + + /** + * Register an Element to the scheduler. + */ + void RegisterElement(const Element &x); + + /** + * Finalize the registration. + */ + void FinalizeRegistration(); + + /** + * Mark this should schedule after another. + * + * @param b + * @param level + */ + Scheduler &After(const Element &a, const Element &b, int level); + /** + * Mark this should schedule before another. + * @param b + * @param level + */ + Scheduler &Before(const Element &a, const Element &b, int level); + + /** + * Build and create schedule. + */ + std::unordered_map BuildSchedule() const; + + private: + /** + * The polyhedral schedule, any schedule is performed on it. + * We use the time-space map to record the schedule infomation, the format is borrowed from Tiramisu project: + * [redundant, + * + */ + int space_size_{}; + //! Tell if the element registration is finalized. + bool registration_finalized_{false}; + //! map from Schedule id to time schedule. + std::unordered_map schedule_; + //! The graph constructed from the dependency and level. There should be only one element which doesn't has + //! dependency and that is the start point. + std::unordered_map depend_flow_graph_; +}; + +} // namespace poly +} // namespace cinn