[Mlir-commits] [mlir] [MLIR] Rerun control for actions in execution context (PR #209197)
Kigyosi Alexandru
llvmlistbot at llvm.org
Tue Jul 21 01:08:33 PDT 2026
https://github.com/akigyosi updated https://github.com/llvm/llvm-project/pull/209197
>From 409df4d8fa956710f8d4c62a6964abefa6497829 Mon Sep 17 00:00:00 2001
From: Alexandru Kigyosi <alexandru.kigyosi at intel.com>
Date: Mon, 13 Jul 2026 14:33:57 +0000
Subject: [PATCH 1/3] rerun control for actions
---
mlir/include/mlir/Debug/ExecutionContext.h | 6 ++++-
mlir/lib/Debug/ExecutionContext.cpp | 22 ++++++++++++++-
mlir/unittests/Debug/ExecutionContextTest.cpp | 27 +++++++++++++++++++
3 files changed, 53 insertions(+), 2 deletions(-)
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
>From 47e9a4cf1478ae197aa12b4f96413e6a6af7b47d Mon Sep 17 00:00:00 2001
From: Alexandru Kigyosi <alexandru.kigyosi at intel.com>
Date: Mon, 13 Jul 2026 14:52:37 +0000
Subject: [PATCH 2/3] apply clang format
---
mlir/include/mlir/Debug/ExecutionContext.h | 9 ++++++++-
1 file changed, 8 insertions(+), 1 deletion(-)
diff --git a/mlir/include/mlir/Debug/ExecutionContext.h b/mlir/include/mlir/Debug/ExecutionContext.h
index e42d61087f22d..93c26e0ce5f82 100644
--- a/mlir/include/mlir/Debug/ExecutionContext.h
+++ b/mlir/include/mlir/Debug/ExecutionContext.h
@@ -66,7 +66,14 @@ class ExecutionContext {
/// reach the parent/enclosing operation. If there are no enclosing
/// operation, the execution continues without stopping.
/// - Rerun: The action gets executed again immediately after it finishes.
- enum Control { Apply = 1, Skip = 2, Step = 3, Next = 4, Finish = 5, Rerun = 6 };
+ 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.
>From 5c211435c096f16a6883f0f581feca09731c3ce4 Mon Sep 17 00:00:00 2001
From: Alexandru Kigyosi <alexandru.kigyosi at intel.com>
Date: Tue, 21 Jul 2026 08:07:15 +0000
Subject: [PATCH 3/3] add nested action test for rerun
---
mlir/unittests/Debug/ExecutionContextTest.cpp | 40 +++++++++++++++++++
1 file changed, 40 insertions(+)
diff --git a/mlir/unittests/Debug/ExecutionContextTest.cpp b/mlir/unittests/Debug/ExecutionContextTest.cpp
index 16c785b6f826e..1c442d4aa1acc 100644
--- a/mlir/unittests/Debug/ExecutionContextTest.cpp
+++ b/mlir/unittests/Debug/ExecutionContextTest.cpp
@@ -376,4 +376,44 @@ TEST(ExecutionContext, RerunRequestedFromPostActionCallback) {
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
More information about the Mlir-commits
mailing list