[Mlir-commits] [mlir] ActionHandler getter (PR #197230)
Kigyosi Alexandru
llvmlistbot at llvm.org
Wed May 13 02:45:11 PDT 2026
https://github.com/akigyosi updated https://github.com/llvm/llvm-project/pull/197230
>From 9c520fdf6896ec7371ab58c491e40ef9248a774e Mon Sep 17 00:00:00 2001
From: Alexandru Kigyosi <alexandru.kigyosi at intel.com>
Date: Tue, 12 May 2026 16:01:12 +0000
Subject: [PATCH 1/2] ActionHandler getter
---
mlir/include/mlir/IR/MLIRContext.h | 4 +
mlir/lib/IR/MLIRContext.cpp | 2 +
mlir/unittests/Debug/ActionHandlerTest.cpp | 111 +++++++++++++++++++++
mlir/unittests/Debug/CMakeLists.txt | 1 +
4 files changed, 118 insertions(+)
create mode 100644 mlir/unittests/Debug/ActionHandlerTest.cpp
diff --git a/mlir/include/mlir/IR/MLIRContext.h b/mlir/include/mlir/IR/MLIRContext.h
index 9690029256474..3f4844f5d040d 100644
--- a/mlir/include/mlir/IR/MLIRContext.h
+++ b/mlir/include/mlir/IR/MLIRContext.h
@@ -267,6 +267,10 @@ class MLIRContext {
/// context. A nullptr handler can be set to disable a previously set handler.
void registerActionHandler(HandlerTy handler);
+ /// Return a copy of the currently registered action handler. Its target can
+ /// be used to gain access to the handler's state, if any.
+ HandlerTy getActionHandler();
+
/// Return true if a valid ActionHandler is set.
bool hasActionHandler();
diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index 7b666d11a4a89..132f03b2aeaf1 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/mlir/lib/IR/MLIRContext.cpp
@@ -380,6 +380,8 @@ void MLIRContext::registerActionHandler(HandlerTy handler) {
getImpl().actionHandler = std::move(handler);
}
+MLIRContext::HandlerTy MLIRContext::getActionHandler() { return getImpl().actionHandler; }
+
/// Dispatch the provided action to the handler if any, or just execute it.
void MLIRContext::executeActionInternal(function_ref<void()> actionFn,
const tracing::Action &action) {
diff --git a/mlir/unittests/Debug/ActionHandlerTest.cpp b/mlir/unittests/Debug/ActionHandlerTest.cpp
new file mode 100644
index 0000000000000..87a828d8ee87f
--- /dev/null
+++ b/mlir/unittests/Debug/ActionHandlerTest.cpp
@@ -0,0 +1,111 @@
+//===- ActionHandlerTest.cpp - Debug Action Handler Tests -----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/Action.h"
+#include "mlir/Support/TypeID.h"
+#include "gmock/gmock.h"
+
+#include <gtest/gtest.h>
+
+#include <mlir/IR/MLIRContext.h>
+#include <mlir/IR/Action.h>
+#include <mlir/Support/LLVM.h>
+#include <mlir/Debug/ExecutionContext.h>
+#include <llvm/ADT/StringRef.h>
+
+#include <memory>
+#include <string>
+#include <vector>
+
+using namespace mlir;
+using namespace mlir::tracing;
+
+namespace {
+
+struct DummyAction final : ActionImpl<DummyAction> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(DummyAction)
+ static constexpr StringLiteral tag = "dummy-action";
+};
+
+} // namespace
+
+namespace {
+
+// State class — lives on the heap, shared across all copies of the handler
+struct HandlerState {
+ bool enabled{true};
+};
+
+/// Owner of a shared_ptr to the state
+/// Every copy of the functor points at the same HandlerState object.
+struct StatefulHandler {
+ std::shared_ptr<HandlerState> state;
+
+ void operator()(mlir::function_ref<void()> actionFn,
+ const mlir::tracing::Action& /*action*/) const {
+ if (!state->enabled) {
+ // Skip execution entirely when disabled.
+ return;
+ }
+ actionFn();
+ }
+};
+
+TEST(ActionHandlerSharedState, SingleCopyEnabledState) {
+ mlir::MLIRContext ctx;
+
+ ctx.registerActionHandler(StatefulHandler{std::make_shared<HandlerState>()});
+
+ // Retrieve a copy of the handler.
+ auto handlerCopy = ctx.getActionHandler();
+ ASSERT_TRUE(static_cast<bool>(handlerCopy));
+
+ int executionCount = 0;
+ auto workFn = [&]() { ++executionCount; };
+
+ DummyAction action;
+
+ handlerCopy(workFn, action);
+
+ // Recover the shared_ptr from the handler copy via target<StatefulHandler>().
+ // target<T>() returns a non-null pointer only when the stored callable type
+ // matches T exactly — which is guaranteed here since we registered StatefulHandler.
+ auto* recovered = handlerCopy.target<StatefulHandler>();
+ ASSERT_NE(recovered, nullptr);
+
+ EXPECT_EQ(executionCount, 1);
+ EXPECT_TRUE(recovered->state->enabled == true);
+}
+
+TEST(ActionHandlerSharedState, MultipleCopiesDisabledState) {
+ mlir::MLIRContext ctx;
+
+ ctx.registerActionHandler(StatefulHandler{std::make_shared<HandlerState>()});
+
+ // Recover the shared_ptr and disable the state
+ auto handlerCopy = ctx.getActionHandler();
+ auto* recovered = handlerCopy.target<StatefulHandler>();
+ ASSERT_NE(recovered, nullptr);
+ recovered->state->enabled = false;
+
+ // A second independent copy also sees enabled==false
+ auto handlerCopy2 = ctx.getActionHandler();
+
+ int executionCount = 0;
+ auto workFn = [&]() { ++executionCount; };
+
+ DummyAction action;
+
+ handlerCopy2(workFn, action);
+
+ // workFn was skipped because the handler saw enabled==false
+ EXPECT_EQ(executionCount, 0);
+ EXPECT_TRUE(recovered->state->enabled == false);
+}
+
+} // namespace
diff --git a/mlir/unittests/Debug/CMakeLists.txt b/mlir/unittests/Debug/CMakeLists.txt
index d55282937956a..2ab00102f93b4 100644
--- a/mlir/unittests/Debug/CMakeLists.txt
+++ b/mlir/unittests/Debug/CMakeLists.txt
@@ -1,4 +1,5 @@
add_mlir_unittest(MLIRDebugTests
+ ActionHandlerTest.cpp
DebugCounterTest.cpp
ExecutionContextTest.cpp
FileLineColLocBreakpointManagerTest.cpp
>From e08640996db761fc548a46b39d097896fa18764b Mon Sep 17 00:00:00 2001
From: Alexandru Kigyosi <alexandru.kigyosi at intel.com>
Date: Wed, 13 May 2026 09:47:09 +0000
Subject: [PATCH 2/2] Pass handler by const reference
---
mlir/include/mlir/IR/MLIRContext.h | 4 +-
mlir/lib/IR/MLIRContext.cpp | 2 +-
mlir/unittests/Debug/ActionHandlerTest.cpp | 46 +++++++---------------
3 files changed, 18 insertions(+), 34 deletions(-)
diff --git a/mlir/include/mlir/IR/MLIRContext.h b/mlir/include/mlir/IR/MLIRContext.h
index 3f4844f5d040d..ac08695b7140a 100644
--- a/mlir/include/mlir/IR/MLIRContext.h
+++ b/mlir/include/mlir/IR/MLIRContext.h
@@ -267,9 +267,9 @@ class MLIRContext {
/// context. A nullptr handler can be set to disable a previously set handler.
void registerActionHandler(HandlerTy handler);
- /// Return a copy of the currently registered action handler. Its target can
+ /// Return a reference to the currently registered action handler. Its target can
/// be used to gain access to the handler's state, if any.
- HandlerTy getActionHandler();
+ const HandlerTy &getActionHandler();
/// Return true if a valid ActionHandler is set.
bool hasActionHandler();
diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index 132f03b2aeaf1..ba9d22518dee3 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/mlir/lib/IR/MLIRContext.cpp
@@ -380,7 +380,7 @@ void MLIRContext::registerActionHandler(HandlerTy handler) {
getImpl().actionHandler = std::move(handler);
}
-MLIRContext::HandlerTy MLIRContext::getActionHandler() { return getImpl().actionHandler; }
+const MLIRContext::HandlerTy &MLIRContext::getActionHandler() { return getImpl().actionHandler; }
/// Dispatch the provided action to the handler if any, or just execute it.
void MLIRContext::executeActionInternal(function_ref<void()> actionFn,
diff --git a/mlir/unittests/Debug/ActionHandlerTest.cpp b/mlir/unittests/Debug/ActionHandlerTest.cpp
index 87a828d8ee87f..95f31250c4440 100644
--- a/mlir/unittests/Debug/ActionHandlerTest.cpp
+++ b/mlir/unittests/Debug/ActionHandlerTest.cpp
@@ -7,20 +7,11 @@
//===----------------------------------------------------------------------===//
#include "mlir/IR/Action.h"
-#include "mlir/Support/TypeID.h"
-#include "gmock/gmock.h"
+#include "mlir/IR/MLIRContext.h"
#include <gtest/gtest.h>
-#include <mlir/IR/MLIRContext.h>
-#include <mlir/IR/Action.h>
-#include <mlir/Support/LLVM.h>
-#include <mlir/Debug/ExecutionContext.h>
-#include <llvm/ADT/StringRef.h>
-
#include <memory>
-#include <string>
-#include <vector>
using namespace mlir;
using namespace mlir::tracing;
@@ -30,19 +21,20 @@ namespace {
struct DummyAction final : ActionImpl<DummyAction> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(DummyAction)
static constexpr StringLiteral tag = "dummy-action";
+
+ DummyAction(llvm::ArrayRef<IRUnit> irUnits) {}
};
} // namespace
namespace {
-// State class — lives on the heap, shared across all copies of the handler
+// State class
struct HandlerState {
bool enabled{true};
};
-/// Owner of a shared_ptr to the state
-/// Every copy of the functor points at the same HandlerState object.
+// Owner of a shared_ptr to the state
struct StatefulHandler {
std::shared_ptr<HandlerState> state;
@@ -56,52 +48,44 @@ struct StatefulHandler {
}
};
-TEST(ActionHandlerSharedState, SingleCopyEnabledState) {
+TEST(ActionHandlerSharedState, EnabledState) {
mlir::MLIRContext ctx;
ctx.registerActionHandler(StatefulHandler{std::make_shared<HandlerState>()});
- // Retrieve a copy of the handler.
- auto handlerCopy = ctx.getActionHandler();
- ASSERT_TRUE(static_cast<bool>(handlerCopy));
+ auto handlerRef = ctx.getActionHandler();
+ ASSERT_TRUE(static_cast<bool>(handlerRef));
int executionCount = 0;
auto workFn = [&]() { ++executionCount; };
- DummyAction action;
-
- handlerCopy(workFn, action);
+ ctx.executeAction<DummyAction>(workFn, {});
- // Recover the shared_ptr from the handler copy via target<StatefulHandler>().
+ // Recover the shared_ptr from the handler via target<StatefulHandler>().
// target<T>() returns a non-null pointer only when the stored callable type
// matches T exactly — which is guaranteed here since we registered StatefulHandler.
- auto* recovered = handlerCopy.target<StatefulHandler>();
+ auto* recovered = handlerRef.target<StatefulHandler>();
ASSERT_NE(recovered, nullptr);
EXPECT_EQ(executionCount, 1);
EXPECT_TRUE(recovered->state->enabled == true);
}
-TEST(ActionHandlerSharedState, MultipleCopiesDisabledState) {
+TEST(ActionHandlerSharedState, DisabledState) {
mlir::MLIRContext ctx;
ctx.registerActionHandler(StatefulHandler{std::make_shared<HandlerState>()});
// Recover the shared_ptr and disable the state
- auto handlerCopy = ctx.getActionHandler();
- auto* recovered = handlerCopy.target<StatefulHandler>();
+ auto handlerRef = ctx.getActionHandler();
+ auto* recovered = handlerRef.target<StatefulHandler>();
ASSERT_NE(recovered, nullptr);
recovered->state->enabled = false;
- // A second independent copy also sees enabled==false
- auto handlerCopy2 = ctx.getActionHandler();
-
int executionCount = 0;
auto workFn = [&]() { ++executionCount; };
- DummyAction action;
-
- handlerCopy2(workFn, action);
+ ctx.executeAction<DummyAction>(workFn, {});
// workFn was skipped because the handler saw enabled==false
EXPECT_EQ(executionCount, 0);
More information about the Mlir-commits
mailing list