Added API to print loop structure

fixstars · Jul 26, 2024 · 459bfb0 · 459bfb0
1 parent 9c57b7f
commit 459bfb0
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/include/ion/builder.h b/include/ion/builder.h
@@ -135,6 +135,12 @@ class Builder {
      */
     const std::map<std::string, Halide::JITExtern>& jit_externs() const;
 
+
+    /** Write out the loop nests specified by the schedule for this
+     * Builder's pipeline. Helpful for understanding what a schedule is
+     * doing. */
+    void print_loop_nest();
+
     /**
      * Register disposer hook which will be called from Builder destructor.
      * This is available only for JIT mode.

diff --git a/src/builder.cc b/src/builder.cc
@@ -293,6 +293,12 @@ const std::map<std::string, Halide::JITExtern>& Builder::jit_externs() const {
     return impl_->jit_externs;
 }
 
+void Builder::print_loop_nest() {
+    if (impl_->pipeline.defined()) {
+        impl_->pipeline.print_loop_nest();
+    }
+}
+
 void Builder::register_disposer(Impl *impl, const std::string& bb_id, const std::string& disposer_symbol) {
     log::info("Builder::register_disposer");
     for (const auto& kv : impl->bb_modules) {