[MLIR][Transform][SMT] Introduce transform.smt.constrain_params

rolfmorel · rolfmorel · commit 7de3873dea46 · 2025-09-21T13:13:12.000-07:00
Introduces a SMT Transform-dialect extension so that we can have an op
to express constrains on Transform-dialect params, in particular when
these params are knobs -- see transform.tune.knob -- and can hence can
be seen as symbolic variables. This op allows expressing joint
constraints over multiple params/knobs together.

While the op's semantics are clearly defined, the operational semantics
- i.e. the `apply()` method - for now just defaults to failure. In the
future we should support attaching an implementation so that users can
Bring Your Own Solver and thereby control performance of interpreting
the op. For now the main usage is to walk schedule IR and collect these
constraints so that knobs can be rewritten to constants that satify the
constraints.
diff --git a/mlir/include/mlir/Dialect/Transform/CMakeLists.txt b/mlir/include/mlir/Dialect/Transform/CMakeLists.txt
@@ -4,5 +4,6 @@ add_subdirectory(IR)
 add_subdirectory(IRDLExtension)
 add_subdirectory(LoopExtension)
 add_subdirectory(PDLExtension)
+add_subdirectory(SMTExtension)
 add_subdirectory(Transforms)
 add_subdirectory(TuneExtension)
diff --git a/mlir/include/mlir/Dialect/Transform/SMTExtension/CMakeLists.txt b/mlir/include/mlir/Dialect/Transform/SMTExtension/CMakeLists.txt
@@ -0,0 +1,6 @@
+set(LLVM_TARGET_DEFINITIONS SMTExtensionOps.td)
+mlir_tablegen(SMTExtensionOps.h.inc -gen-op-decls)
+mlir_tablegen(SMTExtensionOps.cpp.inc -gen-op-defs)
+add_public_tablegen_target(MLIRTransformDialectSMTExtensionOpsIncGen)
+
+add_mlir_doc(SMTExtensionOps SMTExtensionOps Dialects/ -gen-op-doc)
diff --git a/mlir/include/mlir/Dialect/Transform/SMTExtension/SMTExtension.h b/mlir/include/mlir/Dialect/Transform/SMTExtension/SMTExtension.h
@@ -0,0 +1,27 @@
+//===- SMTExtension.h - SMT 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_SMTEXTENSION_SMTEXTENSION_H
+#define MLIR_DIALECT_TRANSFORM_SMTEXTENSION_SMTEXTENSION_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"
+
+namespace mlir {
+class DialectRegistry;
+
+namespace transform {
+/// Registers the SMT extension of the Transform dialect in the given registry.
+void registerSMTExtension(DialectRegistry &dialectRegistry);
+} // namespace transform
+} // namespace mlir
+
+#endif // MLIR_DIALECT_TRANSFORM_SMTEXTENSION_SMTEXTENSION_H
diff --git a/mlir/include/mlir/Dialect/Transform/SMTExtension/SMTExtensionOps.h b/mlir/include/mlir/Dialect/Transform/SMTExtension/SMTExtensionOps.h
@@ -0,0 +1,22 @@
+//===- SMTExtensionOps.h - SMT 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_SMTEXTENSION_SMTEXTENSIONOPS_H
+#define MLIR_DIALECT_TRANSFORM_SMTEXTENSION_SMTEXTENSIONOPS_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"
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Transform/SMTExtension/SMTExtensionOps.h.inc"
+
+
+#endif // MLIR_DIALECT_TRANSFORM_SMTEXTENSION_SMTEXTENSIONOPS_H
diff --git a/mlir/include/mlir/Dialect/Transform/SMTExtension/SMTExtensionOps.td b/mlir/include/mlir/Dialect/Transform/SMTExtension/SMTExtensionOps.td
@@ -0,0 +1,52 @@
+//===- SMTExtensionOps.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_SMTEXTENSION_SMTEXTENSIONOPS
+#define MLIR_DIALECT_TRANSFORM_SMTEXTENSION_SMTEXTENSIONOPS
+
+include "mlir/Dialect/Transform/IR/TransformDialect.td"
+include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+
+def ConstrainParamsOp : Op<Transform_Dialect, "smt.constrain_params", [
+  DeclareOpInterfaceMethods<TransformOpInterface>,
+  DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+  NoTerminator
+]> {
+  let cppNamespace = [{ mlir::transform::smt }];
+
+  let summary = "Express contraints on params interpreted as symbolic values";
+  let description = [{
+    Allows expressing constraints on params using the SMT dialect.
+
+    Each Transform dialect param provided as an operand has a corresponding
+    argument of SMT-type in the region. The SMT-Dialect ops in the region use
+    these arguments as operands.
+
+    The semantics of this op is that all the ops in the region together express
+    a constraint on the params-interpreted-as-smt-vars. The op fails in case the
+    expressed constraint is not satisfiable per SMTLIB semantics. Otherwise the
+    op succeeds.
+
+    ---
+
+    TODO: currently the operational semantics per the Transform interpreter is
+    to always fail. The intention is build out support for hooking in your own
+    operational semantics so you can invoke your favourite solver to determine
+    satisfiability of the corresponding constraint problem.
+  }];
+
+  let arguments = (ins Variadic<TransformParamTypeInterface>:$params);
+  let regions = (region SizedRegion<1>:$body);
+  let assemblyFormat =
+      "`(` $params `)` attr-dict `:` type(operands) $body";
+
+  let hasVerifier = 1;
+}
+
+#endif // MLIR_DIALECT_TRANSFORM_SMTEXTENSION_SMTEXTENSIONOPS
diff --git a/mlir/lib/Bindings/Python/DialectSMT.cpp b/mlir/lib/Bindings/Python/DialectSMT.cpp
@@ -41,6 +41,13 @@ static void populateDialectSMTSubmodule(nanobind::module_ &m) {
                 return mlirSMTTypeGetBitVector(context, width);
               },
               "cls"_a, "width"_a, "context"_a = nb::none());
+  auto smtIntType = mlir_type_subclass(m, "IntType", mlirSMTTypeIsAInt)
+                         .def_classmethod(
+                             "get",
+                             [](const nb::object &, MlirContext context) {
+                               return mlirSMTTypeGetInt(context);
+                             },
+                             "cls"_a, "context"_a.none() = nb::none());
 
   auto exportSMTLIB = [](MlirOperation module, bool inlineSingleUseValues,
                          bool indentLetBody) {
diff --git a/mlir/lib/Dialect/Transform/CMakeLists.txt b/mlir/lib/Dialect/Transform/CMakeLists.txt
@@ -4,6 +4,7 @@ add_subdirectory(IR)
 add_subdirectory(IRDLExtension)
 add_subdirectory(LoopExtension)
 add_subdirectory(PDLExtension)
+add_subdirectory(SMTExtension)
 add_subdirectory(Transforms)
 add_subdirectory(TuneExtension)
 add_subdirectory(Utils)
diff --git a/mlir/lib/Dialect/Transform/SMTExtension/CMakeLists.txt b/mlir/lib/Dialect/Transform/SMTExtension/CMakeLists.txt
@@ -0,0 +1,12 @@
+add_mlir_dialect_library(MLIRTransformSMTExtension
+  SMTExtension.cpp
+  SMTExtensionOps.cpp
+
+  DEPENDS
+  MLIRTransformDialectSMTExtensionOpsIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRTransformDialect
+  MLIRSMT
+)
diff --git a/mlir/lib/Dialect/Transform/SMTExtension/SMTExtension.cpp b/mlir/lib/Dialect/Transform/SMTExtension/SMTExtension.cpp
@@ -0,0 +1,35 @@
+//===- SMTExtension.cpp - SMT 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/SMTExtension/SMTExtension.h"
+#include "mlir/Dialect/Transform/SMTExtension/SMTExtensionOps.h"
+#include "mlir/IR/DialectRegistry.h"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// Transform op registration
+//===----------------------------------------------------------------------===//
+
+namespace {
+class SMTExtension : public transform::TransformDialectExtension<SMTExtension> {
+public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SMTExtension)
+
+  SMTExtension() {
+    registerTransformOps<
+#define GET_OP_LIST
+#include "mlir/Dialect/Transform/SMTExtension/SMTExtensionOps.cpp.inc"
+        >();
+  }
+};
+} // namespace
+
+void mlir::transform::registerSMTExtension(DialectRegistry &dialectRegistry) {
+  dialectRegistry.addExtensions<SMTExtension>();
+}
diff --git a/mlir/lib/Dialect/Transform/SMTExtension/SMTExtensionOps.cpp b/mlir/lib/Dialect/Transform/SMTExtension/SMTExtensionOps.cpp
@@ -0,0 +1,55 @@
+//===- SMTExtensionOps.cpp - SMT 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/SMTExtension/SMTExtensionOps.h"
+#include "mlir/Dialect/SMT/IR/SMTDialect.h"
+#include "mlir/Dialect/Transform/IR/TransformOps.h"
+#include "mlir/Dialect/Transform/SMTExtension/SMTExtension.h"
+
+using namespace mlir;
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Transform/SMTExtension/SMTExtensionOps.cpp.inc"
+
+//===----------------------------------------------------------------------===//
+// ConstrainParamsOp
+//===----------------------------------------------------------------------===//
+
+void transform::smt::ConstrainParamsOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  onlyReadsHandle(getParamsMutable(), effects);
+}
+
+DiagnosedSilenceableFailure
+transform::smt::ConstrainParamsOp::apply(transform::TransformRewriter &rewriter,
+                                         transform::TransformResults &results,
+                                         transform::TransformState &state) {
+  // TODO: Proper operational semantics are to chuck the SMT problem in the body
+  //       to a SMT solver with the arguments of the body constrained to the
+  //       values passed into the op. Success or failure is then determined by
+  //       the solver's result.
+  //       One way to support this is to just promise the TransformOpInterface
+  //       and allow for users to attach their own implementation, which would,
+  //       e.g., translate the ops to SMTLIB and hand that over to the user's
+  //       favourite solver. This requires changes to the dialect's verifier.
+  return emitDefiniteFailure() << "op does not have interpreted semantics yet";
+}
+
+LogicalResult transform::smt::ConstrainParamsOp::verify() {
+  if (getOperands().size() != getBody().getNumArguments())
+    return emitOpError(
+        "must have the same number of block arguments as operands");
+
+  for (auto &op : getBody().getOps()) {
+    if (!isa<mlir::smt::SMTDialect>(op.getDialect()))
+      return emitOpError(
+          "ops contained in region should belong to SMT-dialect");
+  }
+
+  return success();
+}
diff --git a/mlir/lib/RegisterAllExtensions.cpp b/mlir/lib/RegisterAllExtensions.cpp
@@ -53,6 +53,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/SMTExtension/SMTExtension.h"
 #include "mlir/Dialect/Transform/TuneExtension/TuneExtension.h"
 #include "mlir/Dialect/Vector/TransformOps/VectorTransformOps.h"
 #include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
@@ -108,6 +109,7 @@ void mlir::registerAllExtensions(DialectRegistry &registry) {
   transform::registerIRDLExtension(registry);
   transform::registerLoopExtension(registry);
   transform::registerPDLExtension(registry);
+  transform::registerSMTExtension(registry);
   transform::registerTuneExtension(registry);
   vector::registerTransformDialectExtension(registry);
   arm_neon::registerTransformDialectExtension(registry);
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
@@ -171,6 +171,15 @@ ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   DIALECT_NAME transform
   EXTENSION_NAME transform_pdl_extension)
 
+declare_mlir_dialect_extension_python_bindings(
+ADD_TO_PARENT MLIRPythonSources.Dialects
+ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/TransformSMTExtensionOps.td
+  SOURCES
+    dialects/transform/smt.py
+  DIALECT_NAME transform
+  EXTENSION_NAME transform_smt_extension)
+
 declare_mlir_dialect_extension_python_bindings(
 ADD_TO_PARENT MLIRPythonSources.Dialects
 ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
diff --git a/mlir/python/mlir/dialects/TransformSMTExtensionOps.td b/mlir/python/mlir/dialects/TransformSMTExtensionOps.td
@@ -0,0 +1,19 @@
+//===-- TransformSMTExtensionOps.td - Binding entry point --*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Entry point of the generated Python bindings for the SMT extension of the
+// Transform dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef PYTHON_BINDINGS_TRANSFORM_SMT_EXTENSION_OPS
+#define PYTHON_BINDINGS_TRANSFORM_SMT_EXTENSION_OPS
+
+include "mlir/Dialect/Transform/SMTExtension/SMTExtensionOps.td"
+
+#endif // PYTHON_BINDINGS_TRANSFORM_SMT_EXTENSION_OPS
diff --git a/mlir/python/mlir/dialects/smt.py b/mlir/python/mlir/dialects/smt.py
@@ -3,6 +3,7 @@
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from ._smt_ops_gen import *
+from ._smt_enum_gen import *
 
 from .._mlir_libs._mlirDialectsSMT import *
 from ..extras.meta import region_op
diff --git a/mlir/python/mlir/dialects/transform/smt.py b/mlir/python/mlir/dialects/transform/smt.py
@@ -0,0 +1,36 @@
+#  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
+
+from typing import Sequence
+
+from ...ir import Type, Block
+from .._transform_smt_extension_ops_gen import *
+from .._transform_smt_extension_ops_gen import _Dialect
+from ...dialects import transform
+
+try:
+    from .._ods_common import _cext as _ods_cext
+except ImportError as e:
+    raise RuntimeError("Error loading imports from extension module") from e
+
+@_ods_cext.register_operation(_Dialect, replace=True)
+class ConstrainParamsOp(ConstrainParamsOp):
+    def __init__(
+        self,
+        params: Sequence[transform.AnyParamType],
+        arg_types: Sequence[Type],
+        loc=None,
+        ip=None,
+    ):
+        assert len(params) == len(arg_types)
+        super().__init__(
+            params,
+            loc=loc,
+            ip=ip,
+        )
+        self.regions[0].blocks.append(*arg_types)
+
+    @property
+    def body(self) -> Block:
+        return self.regions[0].blocks[0]
diff --git a/mlir/test/Dialect/Transform/test-smt-extension-invalid.mlir b/mlir/test/Dialect/Transform/test-smt-extension-invalid.mlir
@@ -0,0 +1,32 @@
+// RUN: mlir-opt %s --transform-interpreter --split-input-file --verify-diagnostics
+
+// CHECK-LABEL: @constraint_not_using_smt_ops
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @constraint_not_using_smt_ops(%arg0: !transform.any_op {transform.readonly}) {
+    %param_as_param = transform.param.constant 42 -> !transform.param<i64>
+    // expected-error@below {{ops contained in region should belong to SMT-dialect}}
+    transform.smt.constrain_params(%param_as_param) : !transform.param<i64> {
+      ^bb0(%param_as_smt_var: !smt.int):
+      %c4 = arith.constant 4 : i32
+      // This is the kind of thing one might think works:
+      //arith.remsi %param_as_smt_var, %c4 : i32
+    }
+    transform.yield
+  }
+}
+
+// -----
+
+// CHECK-LABEL: @operands_not_one_to_one_with_vars
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @operands_not_one_to_one_with_vars(%arg0: !transform.any_op {transform.readonly}) {
+    %param_as_param = transform.param.constant 42 -> !transform.param<i64>
+    // expected-error@below {{must have the same number of block arguments as operands}}
+    transform.smt.constrain_params(%param_as_param) : !transform.param<i64> {
+      ^bb0(%param_as_smt_var: !smt.int, %param_as_another_smt_var: !smt.int):
+      // This is the kind of thing one might think works:
+      //arith.remsi %param_as_smt_var, %c4 : i32
+    }
+    transform.yield
+  }
+}
diff --git a/mlir/test/Dialect/Transform/test-smt-extension.mlir b/mlir/test/Dialect/Transform/test-smt-extension.mlir
diff --git a/mlir/test/python/dialects/transform_smt_ext.py b/mlir/test/python/dialects/transform_smt_ext.py
