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

[llvmlistbot at llvm.org]

Author: Kigyosi Alexandru
Date: 2026-07-21T11:53:17+02:00
New Revision: d8d1d444a2a69293292f10f19691da05207630f7

URL: https://github.com/llvm/llvm-project/commit/d8d1d444a2a69293292f10f19691da05207630f7
DIFF: https://github.com/llvm/llvm-project/commit/d8d1d444a2a69293292f10f19691da05207630f7.diff

LOG: [MLIR] Rerun control for actions in execution context (#209197)

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.

Added: 
    

Modified: 
    mlir/include/mlir/Debug/ExecutionContext.h
    mlir/lib/Debug/ExecutionContext.cpp
    mlir/unittests/Debug/ExecutionContextTest.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Debug/ExecutionContext.h b/mlir/include/mlir/Debug/ExecutionContext.h
index fbad04b9a2f1d..93c26e0ce5f82 100644
--- a/mlir/include/mlir/Debug/ExecutionContext.h
+++ b/mlir/include/mlir/Debug/ExecutionContext.h
@@ -65,7 +65,15 @@ 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 +142,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..1c442d4aa1acc 100644
--- a/mlir/unittests/Debug/ExecutionContextTest.cpp
+++ b/mlir/unittests/Debug/ExecutionContextTest.cpp
@@ -349,4 +349,71 @@ 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);
+}
+
+TEST(ExecutionContext, RerunStackWithNestedActions) {
+  // Request rerun at an outer depth, then request rerun again at an inner
+  // depth before the outer rerun is consumed.
+  int debuggerHits = 0;
+  int otherHits = 0;
+  int debuggerExecutions = 0;
+  int otherExecutions = 0;
+
+  auto onBreakpoint = [&](const ActionActiveStack *backtrace) {
+    StringRef tag = backtrace->getAction().getTag();
+    if (tag == DebuggerAction::tag)
+      return ++debuggerHits == 1 ? ExecutionContext::Rerun
+                                 : ExecutionContext::Apply;
+    if (tag == OtherAction::tag)
+      return ++otherHits == 1 ? ExecutionContext::Rerun
+                              : ExecutionContext::Apply;
+    ADD_FAILURE();
+    return ExecutionContext::Apply;
+  };
+
+  TagBreakpointManager simpleManager;
+  ExecutionContext executionCtx(onBreakpoint);
+  executionCtx.addBreakpointManager(&simpleManager);
+  simpleManager.addBreakpoint(DebuggerAction::tag);
+  simpleManager.addBreakpoint(OtherAction::tag);
+
+  auto nested = [&]() { ++otherExecutions; };
+  auto original = [&]() {
+    ++debuggerExecutions;
+    executionCtx(nested, OtherAction{});
+  };
+
+  executionCtx(original, DebuggerAction{});
+
+  EXPECT_EQ(debuggerHits, 2);
+  EXPECT_EQ(otherHits, 3);
+  EXPECT_EQ(debuggerExecutions, 2);
+  EXPECT_EQ(otherExecutions, 3);
+}
 } // namespace