From 611b635df3f0e99ea0f8ac4d3722adbdbad824f0 Mon Sep 17 00:00:00 2001
From: Eric Lunderberg <elunderberg@octoml.ai>
Date: Fri, 24 Mar 2023 09:18:58 -0500
Subject: [PATCH] [TVMScript] Improved error message for unexpected top frame

Previously, the error messasge from `FindPrimFuncFrame`,
`FindBlockFrame`, and `FindIfFrame` stated that they could not find
the requested frame when the top-most frame did not match the
requested type.  This error message could be misinterpreted by a user
as stating that the frame didn't exist at all.

This commit updates the error message to distinguish between the case
of a missing frame (e.g. `T.reads()` occurring outside of any
`T.block()` frame) and a frame not being top-most (e.g. `T.reads()`
occurring inside a `T.block()`, but inside an `if` conditional instead
of the top of the block).
---
 src/script/ir_builder/tir/utils.h | 26 ++++++++++++++++++++++----
 1 file changed, 22 insertions(+), 4 deletions(-)

diff --git a/src/script/ir_builder/tir/utils.h b/src/script/ir_builder/tir/utils.h
index 485757063867..7ccc132fa1fe 100644
--- a/src/script/ir_builder/tir/utils.h
+++ b/src/script/ir_builder/tir/utils.h
@@ -69,9 +69,15 @@ inline tvm::tir::Stmt AsStmt(const Array<tvm::tir::Stmt>& stmt) {
 inline PrimFuncFrame FindPrimFuncFrame(const String& method) {
   if (Optional<PrimFuncFrame> frame = IRBuilder::Current()->GetLastFrame<PrimFuncFrame>()) {
     return frame.value();
+  } else if (Optional<PrimFuncFrame> frame = IRBuilder::Current()->FindFrame<PrimFuncFrame>()) {
+    LOG(FATAL) << "ValueError: " << method << " must be called at the top of a PrimFunc.  "
+               << "While " << method << " did occur within the PrimFunc \"" << frame.value()->name
+               << "\", other frames (e.g. block/if/else/let) had been introduced since the "
+               << "PrimFunc's frame";
+  } else {
+    LOG(FATAL) << "ValueError: " << method << " must be called at the top of a PrimFunc, "
+               << "but " << method << " occurred outside of any T.prim_func() frame";
   }
-  LOG(FATAL) << "ValueError: PrimFunc frame not find. Please ensure '" << method
-             << "' is called under T.prim_func()";
   throw;
 }
 
@@ -83,9 +89,15 @@ inline PrimFuncFrame FindPrimFuncFrame(const String& method) {
 inline BlockFrame FindBlockFrame(const String& method) {
   if (Optional<BlockFrame> frame = IRBuilder::Current()->GetLastFrame<BlockFrame>()) {
     return frame.value();
+  } else if (Optional<BlockFrame> frame = IRBuilder::Current()->FindFrame<BlockFrame>()) {
+    LOG(FATAL) << "ValueError: " << method << " must be called at the top of a T.block().  "
+               << "While " << method << " did occur within the block \"" << frame.value()->name
+               << "\", other frames (e.g. if/else/let) had been introduced since the T.block(\""
+               << frame.value()->name << "\") frame";
+  } else {
+    LOG(FATAL) << "ValueError: " << method << " must be called at the top of a T.block(), "
+               << "but " << method << " occurred outside of any T.block() frame";
   }
-  LOG(FATAL) << "ValueError: Block frame not find. Please ensure '" << method
-             << "' is called under T.block()";
   throw;
 }
 
@@ -97,6 +109,12 @@ inline BlockFrame FindBlockFrame(const String& method) {
 inline IfFrame FindIfFrame(const String& method) {
   if (Optional<IfFrame> frame = IRBuilder::Current()->GetLastFrame<IfFrame>()) {
     return frame.value();
+  } else if (Optional<IfFrame> frame = IRBuilder::Current()->FindFrame<IfFrame>()) {
+    LOG(FATAL) << "ValueError: " << method << " must be called at the top of a T.if_().  "
+               << "While " << method << " did occur within the conditional based on ("
+               << frame.value()->condition
+               << "), other frames (e.g. if/else/let) had been introduced since the "
+               << "IfThenElse frame";
   } else {
     LOG(FATAL) << "ValueError: IfThenElse frame not find. Please ensure '" << method
                << "' is called under T.if_()";