[MLIR][Transform][SMT] Introduce transform.smt.constrain_params (#159450)

Introduces a Transform-dialect SMT-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 be
seen as symbolic variables. This op allows expressing joint constraints
over multiple params/knobs together.

While the op's semantics are clearly defined, per SMTLIB, the interpreted
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
satisfy the constraints.

Author: rolfmorel

mlir/lib/Bindings/Python/DialectSMT.cpp

@@ -26,21 +26,26 @@ using namespace mlir::python::nanobind_adaptors;
 
 static void populateDialectSMTSubmodule(nanobind::module_ &m) {
 
-  auto smtBoolType = mlir_type_subclass(m, "BoolType", mlirSMTTypeIsABool)
-                         .def_classmethod(
-                             "get",
-                             [](const nb::object &, MlirContext context) {
-                               return mlirSMTTypeGetBool(context);
-                             },
-                             "cls"_a, "context"_a = nb::none());
+  auto smtBoolType =
+      mlir_type_subclass(m, "BoolType", mlirSMTTypeIsABool)
+          .def_staticmethod(
+              "get",
+              [](MlirContext context) { return mlirSMTTypeGetBool(context); },
+              "context"_a = nb::none());
   auto smtBitVectorType =
       mlir_type_subclass(m, "BitVectorType", mlirSMTTypeIsABitVector)
-          .def_classmethod(
+          .def_staticmethod(
               "get",
-              [](const nb::object &, int32_t width, MlirContext context) {
+              [](int32_t width, MlirContext context) {
                 return mlirSMTTypeGetBitVector(context, width);
               },
-              "cls"_a, "width"_a, "context"_a = nb::none());
+              "width"_a, "context"_a = nb::none());
+  auto smtIntType =
+      mlir_type_subclass(m, "IntType", mlirSMTTypeIsAInt)
+          .def_staticmethod(
+              "get",
+              [](MlirContext context) { return mlirSMTTypeGetInt(context); },
+              "context"_a = nb::none());
 
   auto exportSMTLIB = [](MlirOperation module, bool inlineSingleUseValues,
                          bool indentLetBody) {

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"

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

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

mlir/python/mlir/dialects/transform/smt.py

@@ -0,0 +1,38 @@
+#  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,
+    ):
+        if len(params) != len(arg_types):
+            raise ValueError(f"{params=} not same length as {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]