[Mlir-commits] [mlir] [MLIR][Transform] Introduce `transform.tune.select` op (PR #146732)
Rolf Morel
llvmlistbot at llvm.org
Wed Jul 2 08:50:09 PDT 2025
https://github.com/rolfmorel created https://github.com/llvm/llvm-project/pull/146732
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.
>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] [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();
+}
More information about the Mlir-commits
mailing list