[Mlir-commits] [mlir] [MLIR] Rerun control for actions in execution context (PR #209197)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Kigyosi Alexandru (akigyosi)

<details>
<summary>Changes</summary>

Adds a Rerun control value to ExecutionContext::Control that allows the re-execution of the current action immediately after it completes, without restarting the full compilation pipeline. This is analogous to GDB's ability to restart execution from a breakpoint. When the callback returns Rerun, the action is executed normally, then re-dispatched through the full ExecutionContext::operator() pipeline, including breakpoint matching, so the user gets a fresh opportunity to inspect or control the re-execution. As a practical usage example, a breakpoint + an observer can be added, to save and restore IR between runs to check if each run produces the same IR or something different each time. A depth-keyed structure is used, so rerun requests survive nested action dispatch and are consumed by the correct stack frame.

---
Full diff: https://github.com/llvm/llvm-project/pull/209197.diff


3 Files Affected:

- (modified) mlir/include/mlir/Debug/ExecutionContext.h (+5-1) 
- (modified) mlir/lib/Debug/ExecutionContext.cpp (+21-1) 
- (modified) mlir/unittests/Debug/ExecutionContextTest.cpp (+27) 


``````````diff
diff --git a/mlir/include/mlir/Debug/ExecutionContext.h b/mlir/include/mlir/Debug/ExecutionContext.h
index fbad04b9a2f1d..e42d61087f22d 100644
--- a/mlir/include/mlir/Debug/ExecutionContext.h
+++ b/mlir/include/mlir/Debug/ExecutionContext.h
@@ -65,7 +65,8 @@ class ExecutionContext {
   /// - Finish: The action is executed and the execution is paused only when we
   ///           reach the parent/enclosing operation. If there are no enclosing
   ///           operation, the execution continues without stopping.
-  enum Control { Apply = 1, Skip = 2, Step = 3, Next = 4, Finish = 5 };
+  /// - Rerun: The action gets executed again immediately after it finishes.
+  enum Control { Apply = 1, Skip = 2, Step = 3, Next = 4, Finish = 5, Rerun = 6 };
 
   /// The type of the callback that is used to control the execution.
   /// The callback is passed the current action.
@@ -134,6 +135,9 @@ class ExecutionContext {
 
   /// The list of managers that are queried for breakpoints.
   SmallVector<BreakpointManager *> breakpoints;
+
+  /// Stack of depths of actions that should be rerun.
+  SmallVector<int> rerunControlStack;
 };
 
 } // namespace tracing
diff --git a/mlir/lib/Debug/ExecutionContext.cpp b/mlir/lib/Debug/ExecutionContext.cpp
index 7fc5c165d391e..cc72bac549ab0 100644
--- a/mlir/lib/Debug/ExecutionContext.cpp
+++ b/mlir/lib/Debug/ExecutionContext.cpp
@@ -84,10 +84,25 @@ void ExecutionContext::operator()(llvm::function_ref<void()> transform,
     case ExecutionContext::Finish:
       depthToBreak = depth - 1;
       return true;
+    case ExecutionContext::Rerun:
+      depthToBreak = std::nullopt;
+      rerunControlStack.push_back(depth);
+      return true;
     }
     llvm::report_fatal_error("Unknown control request");
   };
 
+  auto rerunCurrentActionIfRequested = [&]() -> bool {
+    if (!rerunControlStack.empty() && rerunControlStack.back() == depth) {
+      // If the user requested to rerun this action, we do it here.
+      rerunControlStack.pop_back();
+      actionStack = info.getParent();
+      (*this)(transform, action);
+      return true;
+    }
+    return false;
+  };
+
   // Try to find a breakpoint that would hit on this action.
   // Right now there is no way to collect them all, we stop at the first one.
   for (auto *breakpointManager : breakpoints) {
@@ -116,6 +131,11 @@ void ExecutionContext::operator()(llvm::function_ref<void()> transform,
       observer->afterExecute(actionStack);
   }
 
-  if (depthToBreak && depth <= depthToBreak)
+  if (rerunCurrentActionIfRequested())
+    return;
+
+  if (depthToBreak && depth <= depthToBreak) {
     handleUserInput();
+    rerunCurrentActionIfRequested();
+  }
 }
diff --git a/mlir/unittests/Debug/ExecutionContextTest.cpp b/mlir/unittests/Debug/ExecutionContextTest.cpp
index 642adff51002a..16c785b6f826e 100644
--- a/mlir/unittests/Debug/ExecutionContextTest.cpp
+++ b/mlir/unittests/Debug/ExecutionContextTest.cpp
@@ -349,4 +349,31 @@ TEST(ExecutionContext, EnableDisableBreakpointOnCallback) {
   executionCtx(original, DebuggerAction{});
   EXPECT_EQ(counter, 4);
 }
+
+TEST(ExecutionContext, RerunRequestedFromPostActionCallback) {
+  // If rerun is requested from the post-action callback path,
+  // the same action must be rerun immediately.
+  std::vector<ExecutionContext::Control> controlSequence = {
+      ExecutionContext::Next, ExecutionContext::Rerun, ExecutionContext::Apply};
+  int idx = 0;
+  int callbackCounter = 0;
+  int executionCounter = 0;
+
+  auto onBreakpoint = [&](const ActionActiveStack *backtrace) {
+    ++callbackCounter;
+    EXPECT_EQ(backtrace->getAction().getTag(), DebuggerAction::tag);
+    return controlSequence[idx++];
+  };
+
+  TagBreakpointManager simpleManager;
+  ExecutionContext executionCtx(onBreakpoint);
+  executionCtx.addBreakpointManager(&simpleManager);
+  simpleManager.addBreakpoint(DebuggerAction::tag);
+
+  auto callback = [&]() { ++executionCounter; };
+
+  executionCtx(callback, DebuggerAction{});
+  EXPECT_EQ(callbackCounter, 3);
+  EXPECT_EQ(executionCounter, 2);
+}
 } // namespace

``````````

</details>


https://github.com/llvm/llvm-project/pull/209197