From 459bfb05869cd5b1a154f2111ba36eec4304709c Mon Sep 17 00:00:00 2001
From: Takuro Iizuka <iitaku@gmail.com>
Date: Fri, 26 Jul 2024 13:42:16 -0700
Subject: [PATCH] Added API to print loop structure

---
 include/ion/builder.h | 6 ++++++
 src/builder.cc        | 6 ++++++
 2 files changed, 12 insertions(+)

diff --git a/include/ion/builder.h b/include/ion/builder.h
index 074c54a9..94c6bd94 100644
--- 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
index 3227194f..d96eabf9 100644
--- 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) {