[Mlir-commits] [mlir] [MLIR][Transform] Introduce `transform.tune.select` op (PR #146732)
Rolf Morel
llvmlistbot at llvm.org
Thu Jul 3 06:54:15 PDT 2025
https://github.com/rolfmorel updated https://github.com/llvm/llvm-project/pull/146732
>From 25ab701d44721727a56ab5de8ca13021abad0db3 Mon Sep 17 00:00:00 2001
From: Rolf Morel <rolf.morel at intel.com>
Date: Wed, 2 Jul 2025 08:35:02 -0700
Subject: [PATCH 1/2] [MLIR][Transform] Introduce `transform.tune.select` op
A new op to represent that an attribute is to be chosen from a set of
alternatives and that this choice is made available as a
`!transform.param`. When a `selected` argument is provided, the op's
`apply()` semantics is that of just making this selected attribute
available as the result. When `selected` is not provided, `apply()`
complains that nothing has resolved the non-determinism that the op is
representing.
---
.../mlir/Dialect/Transform/CMakeLists.txt | 1 +
.../Transform/TuneExtension/CMakeLists.txt | 6 ++
.../Transform/TuneExtension/TuneExtension.h | 21 +++++++
.../TuneExtension/TuneExtensionOps.h | 22 +++++++
.../TuneExtension/TuneExtensionOps.td | 36 +++++++++++
mlir/lib/Dialect/Transform/CMakeLists.txt | 1 +
.../Transform/TuneExtension/CMakeLists.txt | 12 ++++
.../Transform/TuneExtension/TuneExtension.cpp | 34 +++++++++++
.../TuneExtension/TuneExtensionOps.cpp | 61 +++++++++++++++++++
9 files changed, 194 insertions(+)
create mode 100644 mlir/include/mlir/Dialect/Transform/TuneExtension/CMakeLists.txt
create mode 100644 mlir/include/mlir/Dialect/Transform/TuneExtension/TuneExtension.h
create mode 100644 mlir/include/mlir/Dialect/Transform/TuneExtension/TuneExtensionOps.h
create mode 100644 mlir/include/mlir/Dialect/Transform/TuneExtension/TuneExtensionOps.td
create mode 100644 mlir/lib/Dialect/Transform/TuneExtension/CMakeLists.txt
create mode 100644 mlir/lib/Dialect/Transform/TuneExtension/TuneExtension.cpp
create mode 100644 mlir/lib/Dialect/Transform/TuneExtension/TuneExtensionOps.cpp
diff --git a/mlir/include/mlir/Dialect/Transform/CMakeLists.txt b/mlir/include/mlir/Dialect/Transform/CMakeLists.txt
index b6155b5f573f1..e70479b2a39f2 100644
--- a/mlir/include/mlir/Dialect/Transform/CMakeLists.txt
+++ b/mlir/include/mlir/Dialect/Transform/CMakeLists.txt
@@ -5,3 +5,4 @@ add_subdirectory(IRDLExtension)
add_subdirectory(LoopExtension)
add_subdirectory(PDLExtension)
add_subdirectory(Transforms)
+add_subdirectory(TuneExtension)
diff --git a/mlir/include/mlir/Dialect/Transform/TuneExtension/CMakeLists.txt b/mlir/include/mlir/Dialect/Transform/TuneExtension/CMakeLists.txt
new file mode 100644
index 0000000000000..9afca813afda6
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Transform/TuneExtension/CMakeLists.txt
@@ -0,0 +1,6 @@
+set(LLVM_TARGET_DEFINITIONS TuneExtensionOps.td)
+mlir_tablegen(TuneExtensionOps.h.inc -gen-op-decls)
+mlir_tablegen(TuneExtensionOps.cpp.inc -gen-op-defs)
+add_public_tablegen_target(MLIRTransformDialectTuneExtensionOpsIncGen)
+
+add_mlir_doc(TuneExtensionOps TuneExtensionOps Dialects/ -gen-op-doc)
diff --git a/mlir/include/mlir/Dialect/Transform/TuneExtension/TuneExtension.h b/mlir/include/mlir/Dialect/Transform/TuneExtension/TuneExtension.h
new file mode 100644
index 0000000000000..1453d1754297f
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Transform/TuneExtension/TuneExtension.h
@@ -0,0 +1,21 @@
+//===- TuneExtension.h - Tune extension for Transform dialect ---*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_TRANSFORM_TUNEEXTENSION_TUNEEXTENSION_H
+#define MLIR_DIALECT_TRANSFORM_TUNEEXTENSION_TUNEEXTENSION_H
+
+namespace mlir {
+class DialectRegistry;
+
+namespace transform {
+/// Registers the tune extension of the Transform dialect in the given registry.
+void registerTuneExtension(DialectRegistry &dialectRegistry);
+} // namespace transform
+} // namespace mlir
+
+#endif // MLIR_DIALECT_TRANSFORM_TUNEEXTENSION_TUNEEXTENSION_H
diff --git a/mlir/include/mlir/Dialect/Transform/TuneExtension/TuneExtensionOps.h b/mlir/include/mlir/Dialect/Transform/TuneExtension/TuneExtensionOps.h
new file mode 100644
index 0000000000000..de5bbc61919e9
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Transform/TuneExtension/TuneExtensionOps.h
@@ -0,0 +1,22 @@
+//===- TuneExtensionOps.h - Tune ext. for Transform dialect -----*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_TRANSFORM_TUNEEXTENSION_TUNEEXTENSIONOPS_H
+#define MLIR_DIALECT_TRANSFORM_TUNEEXTENSION_TUNEEXTENSIONOPS_H
+
+#include "mlir/Bytecode/BytecodeOpInterface.h"
+#include "mlir/Dialect/Transform/IR/TransformDialect.h"
+#include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Transform/TuneExtension/TuneExtensionOps.h.inc"
+
+#endif // MLIR_DIALECT_TRANSFORM_TUNEEXTENSION_TUNEEXTENSIONOPS_H
diff --git a/mlir/include/mlir/Dialect/Transform/TuneExtension/TuneExtensionOps.td b/mlir/include/mlir/Dialect/Transform/TuneExtension/TuneExtensionOps.td
new file mode 100644
index 0000000000000..9366ab0ddd240
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Transform/TuneExtension/TuneExtensionOps.td
@@ -0,0 +1,36 @@
+//===- TuneExtensionOps.td - Transform dialect operations --*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_TRANSFORM_TUNEEXTENSION_TUNEEXTENSIONOPS
+#define MLIR_DIALECT_TRANSFORM_TUNEEXTENSION_TUNEEXTENSIONOPS
+
+include "mlir/Dialect/Transform/IR/TransformDialect.td"
+include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/IR/CommonAttrConstraints.td"
+
+def SelectOp : Op<Transform_Dialect, "tune.select", [
+ DeclareOpInterfaceMethods<TransformOpInterface>,
+ DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+]> {
+ let summary = "Non-deterministically select a value from a set of values";
+ let description = [{
+ TODO
+ }];
+ let cppNamespace = [{ mlir::transform::tune }];
+ let hasVerifier = 1;
+
+ let arguments = (ins SymbolRefAttr:$name,
+ AnyAttr:$options,
+ OptionalAttr<AnyAttr>:$selected);
+ let results = (outs TransformParamTypeInterface:$result);
+ let assemblyFormat =
+ "$name (`=` $selected^ `selected`)? `from` $options attr-dict `->` type(results)";
+}
+
+#endif // MLIR_DIALECT_TRANSFORM_TUNEEXTENSION_TUNEEXTENSIONOPS
diff --git a/mlir/lib/Dialect/Transform/CMakeLists.txt b/mlir/lib/Dialect/Transform/CMakeLists.txt
index 0c0d5ebe0c212..6e628353258d6 100644
--- a/mlir/lib/Dialect/Transform/CMakeLists.txt
+++ b/mlir/lib/Dialect/Transform/CMakeLists.txt
@@ -5,4 +5,5 @@ add_subdirectory(IRDLExtension)
add_subdirectory(LoopExtension)
add_subdirectory(PDLExtension)
add_subdirectory(Transforms)
+add_subdirectory(TuneExtension)
add_subdirectory(Utils)
diff --git a/mlir/lib/Dialect/Transform/TuneExtension/CMakeLists.txt b/mlir/lib/Dialect/Transform/TuneExtension/CMakeLists.txt
new file mode 100644
index 0000000000000..ff01d25e57f68
--- /dev/null
+++ b/mlir/lib/Dialect/Transform/TuneExtension/CMakeLists.txt
@@ -0,0 +1,12 @@
+add_mlir_dialect_library(MLIRTransformTuneExtension
+ TuneExtension.cpp
+ TuneExtensionOps.cpp
+
+ DEPENDS
+ MLIRTransformDialectTuneExtensionOpsIncGen
+
+ LINK_LIBS PUBLIC
+ MLIRIR
+ MLIRTransformDialect
+ MLIRTransforms
+)
diff --git a/mlir/lib/Dialect/Transform/TuneExtension/TuneExtension.cpp b/mlir/lib/Dialect/Transform/TuneExtension/TuneExtension.cpp
new file mode 100644
index 0000000000000..c4581db83bb05
--- /dev/null
+++ b/mlir/lib/Dialect/Transform/TuneExtension/TuneExtension.cpp
@@ -0,0 +1,34 @@
+//===- TuneExtension.cpp - Tune extension for the Transform dialect -------===//
+//
+// 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/Dialect/Transform/TuneExtension/TuneExtension.h"
+
+#include "mlir/Dialect/Transform/IR/TransformDialect.h"
+#include "mlir/Dialect/Transform/TuneExtension/TuneExtensionOps.h"
+#include "mlir/IR/DialectRegistry.h"
+
+using namespace mlir;
+
+/// Tune extension of the Transform dialect. This provides "core" transform
+/// operations for loop-like ops.
+class TuneExtension
+ : public transform::TransformDialectExtension<TuneExtension> {
+public:
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TuneExtension)
+
+ void init() {
+ registerTransformOps<
+#define GET_OP_LIST
+#include "mlir/Dialect/Transform/TuneExtension/TuneExtensionOps.cpp.inc"
+ >();
+ }
+};
+
+void mlir::transform::registerTuneExtension(DialectRegistry &dialectRegistry) {
+ dialectRegistry.addExtensions<TuneExtension>();
+}
diff --git a/mlir/lib/Dialect/Transform/TuneExtension/TuneExtensionOps.cpp b/mlir/lib/Dialect/Transform/TuneExtension/TuneExtensionOps.cpp
new file mode 100644
index 0000000000000..401f09eb4b6dc
--- /dev/null
+++ b/mlir/lib/Dialect/Transform/TuneExtension/TuneExtensionOps.cpp
@@ -0,0 +1,61 @@
+//===- TuneExtensionOps.cpp - Tune extension for the Transform dialect ----===//
+//
+// 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/Dialect/Transform/TuneExtension/TuneExtensionOps.h"
+#include "mlir/Dialect/Transform/IR/TransformOps.h"
+#include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"
+
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/PatternMatch.h"
+#include "llvm/Support/Debug.h"
+
+using namespace mlir;
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Transform/TuneExtension/TuneExtensionOps.cpp.inc"
+
+#define DEBUG_TYPE "transform-tune"
+#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "] ")
+
+//===----------------------------------------------------------------------===//
+// SelectOp
+//===----------------------------------------------------------------------===//
+
+void transform::tune::SelectOp::getEffects(
+ SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+ producesHandle(getOperation()->getOpResults(), effects);
+ onlyReadsPayload(effects);
+}
+
+DiagnosedSilenceableFailure
+transform::tune::SelectOp::apply(transform::TransformRewriter &rewriter,
+ transform::TransformResults &results,
+ transform::TransformState &state) {
+ if (getSelected()) {
+ results.setParams(getOperation()->getOpResults()[0], *getSelected());
+ return DiagnosedSilenceableFailure::success();
+ }
+
+ return emitDefiniteFailure() << "non-deterministic choice is only resolved "
+ "through providing a `selected` attr!";
+}
+
+LogicalResult transform::tune::SelectOp::verify() {
+ if (auto selected = getSelected()) {
+ if (auto optionsArray = dyn_cast<ArrayAttr>(getOptions())) {
+ if (!llvm::is_contained(optionsArray, selected))
+ return emitOpError("provided `selected` attribute is not an element of "
+ "`options` array of attributes");
+ } else
+ LLVM_DEBUG(DBGS() << "cannot verify `selected` attribute " << selected
+ << " is an element of `options` attribute "
+ << getOptions());
+ }
+
+ return success();
+}
>From 020562047c425177f6307c0fcf884736371ae2eb Mon Sep 17 00:00:00 2001
From: Rolf Morel <rolf.morel at intel.com>
Date: Thu, 3 Jul 2025 06:45:47 -0700
Subject: [PATCH 2/2] Move to new name, KnobOp, and new syntax and add test
cases and docs
---
.../TuneExtension/TuneExtensionOps.td | 28 +++++++++++++++----
mlir/include/mlir/InitAllExtensions.h | 2 ++
.../Transform/TuneExtension/TuneExtension.cpp | 2 --
.../TuneExtension/TuneExtensionOps.cpp | 17 +++++------
mlir/python/CMakeLists.txt | 9 ++++++
5 files changed, 43 insertions(+), 15 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Transform/TuneExtension/TuneExtensionOps.td b/mlir/include/mlir/Dialect/Transform/TuneExtension/TuneExtensionOps.td
index 9366ab0ddd240..afb67e8fef250 100644
--- a/mlir/include/mlir/Dialect/Transform/TuneExtension/TuneExtensionOps.td
+++ b/mlir/include/mlir/Dialect/Transform/TuneExtension/TuneExtensionOps.td
@@ -12,25 +12,43 @@
include "mlir/Dialect/Transform/IR/TransformDialect.td"
include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/IR/BuiltinAttributes.td"
include "mlir/IR/CommonAttrConstraints.td"
-def SelectOp : Op<Transform_Dialect, "tune.select", [
+def KnobOp : Op<Transform_Dialect, "tune.knob", [
DeclareOpInterfaceMethods<TransformOpInterface>,
DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
]> {
- let summary = "Non-deterministically select a value from a set of values";
+ let summary = "Represents a tunable parameter with a set of options";
+
let description = [{
- TODO
+ Provides a representation for "tunables" within schedules.
+
+ Each op represents a single tunable, which has a `name` and a set
+ of valid `options` described by an attribute. Without a specified
+ `selected` option, this op represents a non-deterministic choice
+ that has yet to be resolved -- as such, the interpreter runtime
+ semantics is to raise a failure.
+
+ The non-deterministic choice is resolved through providing a
+ `selected` attribute. When provided, the interpreter runtime
+ semantics are to return the `selected` attribute as a param through
+ the op's result.
+
+ -----
+
+ In case the `options` attribute is an `ArrayAttr`, the verifier checks that the provided `selected` attribute occurs in `options`.
}];
let cppNamespace = [{ mlir::transform::tune }];
let hasVerifier = 1;
- let arguments = (ins SymbolRefAttr:$name,
+ let arguments = (ins Builtin_StringAttr:$name,
AnyAttr:$options,
OptionalAttr<AnyAttr>:$selected);
let results = (outs TransformParamTypeInterface:$result);
+
let assemblyFormat =
- "$name (`=` $selected^ `selected`)? `from` $options attr-dict `->` type(results)";
+ "`<` $name `>` (`=` $selected^ `from`)? `options` `=` $options attr-dict `->` type(results)";
}
#endif // MLIR_DIALECT_TRANSFORM_TUNEEXTENSION_TUNEEXTENSIONOPS
diff --git a/mlir/include/mlir/InitAllExtensions.h b/mlir/include/mlir/InitAllExtensions.h
index f356b91b1b6c0..0f2d0e45008cc 100644
--- a/mlir/include/mlir/InitAllExtensions.h
+++ b/mlir/include/mlir/InitAllExtensions.h
@@ -52,6 +52,7 @@
#include "mlir/Dialect/Transform/IRDLExtension/IRDLExtension.h"
#include "mlir/Dialect/Transform/LoopExtension/LoopExtension.h"
#include "mlir/Dialect/Transform/PDLExtension/PDLExtension.h"
+#include "mlir/Dialect/Transform/TuneExtension/TuneExtension.h"
#include "mlir/Dialect/Vector/TransformOps/VectorTransformOps.h"
#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
#include "mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h"
@@ -107,6 +108,7 @@ inline void registerAllExtensions(DialectRegistry ®istry) {
transform::registerIRDLExtension(registry);
transform::registerLoopExtension(registry);
transform::registerPDLExtension(registry);
+ transform::registerTuneExtension(registry);
vector::registerTransformDialectExtension(registry);
arm_neon::registerTransformDialectExtension(registry);
arm_sve::registerTransformDialectExtension(registry);
diff --git a/mlir/lib/Dialect/Transform/TuneExtension/TuneExtension.cpp b/mlir/lib/Dialect/Transform/TuneExtension/TuneExtension.cpp
index c4581db83bb05..e18f1e2748540 100644
--- a/mlir/lib/Dialect/Transform/TuneExtension/TuneExtension.cpp
+++ b/mlir/lib/Dialect/Transform/TuneExtension/TuneExtension.cpp
@@ -14,8 +14,6 @@
using namespace mlir;
-/// Tune extension of the Transform dialect. This provides "core" transform
-/// operations for loop-like ops.
class TuneExtension
: public transform::TransformDialectExtension<TuneExtension> {
public:
diff --git a/mlir/lib/Dialect/Transform/TuneExtension/TuneExtensionOps.cpp b/mlir/lib/Dialect/Transform/TuneExtension/TuneExtensionOps.cpp
index 401f09eb4b6dc..6ee0753b3d556 100644
--- a/mlir/lib/Dialect/Transform/TuneExtension/TuneExtensionOps.cpp
+++ b/mlir/lib/Dialect/Transform/TuneExtension/TuneExtensionOps.cpp
@@ -23,29 +23,30 @@ using namespace mlir;
#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "] ")
//===----------------------------------------------------------------------===//
-// SelectOp
+// KnobOp
//===----------------------------------------------------------------------===//
-void transform::tune::SelectOp::getEffects(
+void transform::tune::KnobOp::getEffects(
SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
producesHandle(getOperation()->getOpResults(), effects);
onlyReadsPayload(effects);
}
DiagnosedSilenceableFailure
-transform::tune::SelectOp::apply(transform::TransformRewriter &rewriter,
- transform::TransformResults &results,
- transform::TransformState &state) {
+transform::tune::KnobOp::apply(transform::TransformRewriter &rewriter,
+ transform::TransformResults &results,
+ transform::TransformState &state) {
if (getSelected()) {
results.setParams(getOperation()->getOpResults()[0], *getSelected());
return DiagnosedSilenceableFailure::success();
}
- return emitDefiniteFailure() << "non-deterministic choice is only resolved "
- "through providing a `selected` attr!";
+ return emitDefiniteFailure()
+ << "non-deterministic choice " << getName()
+ << " is only resolved through providing a `selected` attr";
}
-LogicalResult transform::tune::SelectOp::verify() {
+LogicalResult transform::tune::KnobOp::verify() {
if (auto selected = getSelected()) {
if (auto optionsArray = dyn_cast<ArrayAttr>(getOptions())) {
if (!llvm::is_contained(optionsArray, selected))
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index b2daabb2a5957..7a0c95ebb8200 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -180,6 +180,15 @@ ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
DIALECT_NAME transform
EXTENSION_NAME transform_debug_extension)
+declare_mlir_dialect_extension_python_bindings(
+ADD_TO_PARENT MLIRPythonSources.Dialects
+ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+ TD_FILE dialects/TransformTuneExtensionOps.td
+ SOURCES
+ dialects/transform/tune.py
+ DIALECT_NAME transform
+ EXTENSION_NAME transform_tune_extension)
+
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
More information about the Mlir-commits
mailing list