Address reviewer comments

Author: rolfmorel

mlir/lib/Bindings/Python/DialectSMT.cpp

@@ -26,28 +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());
-  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());
+              "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/lib/Dialect/Transform/SMTExtension/SMTExtensionOps.cpp

@@ -29,8 +29,8 @@ 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
+  // TODO: Proper operational semantics are to check the SMT problem in the body
+  //       with 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

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

@@ -24,7 +24,8 @@ def __init__(
         loc=None,
         ip=None,
     ):
-        assert len(params) == len(arg_types)
+        if len(params) != len(arg_types):
+            raise ValueError(f"{params=} not same length as {arg_types=}")
         super().__init__(
             params,
             loc=loc,

mlir/test/Dialect/Transform/test-smt-extension-invalid.mlir

@@ -24,8 +24,6 @@ module attributes {transform.with_named_sequence} {
     // 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
   }

mlir/test/python/dialects/transform_smt_ext.py

@@ -1,21 +1,21 @@
 # RUN: %PYTHON %s | FileCheck %s
 
-from mlir.ir import *
+from mlir import ir
 from mlir.dialects import transform, smt
 from mlir.dialects.transform import smt as transform_smt
 
 
 def run(f):
     print("\nTEST:", f.__name__)
-    with Context(), Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
+    with ir.Context(), ir.Location.unknown():
+        module = ir.Module.create()
+        with ir.InsertionPoint(module.body):
             sequence = transform.SequenceOp(
                 transform.FailurePropagationMode.Propagate,
                 [],
                 transform.AnyOpType.get(),
             )
-            with InsertionPoint(sequence.body):
+            with ir.InsertionPoint(sequence.body):
                 f(sequence.bodyTarget)
                 transform.YieldOp()
         print(module)
@@ -25,7 +25,7 @@ def run(f):
 # CHECK-LABEL: TEST: testConstrainParamsOp
 @run
 def testConstrainParamsOp(target):
-    dummy_value = IntegerAttr.get(IntegerType.get_signless(32), 42)
+    dummy_value = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 42)
     # CHECK: %[[PARAM_AS_PARAM:.*]] = transform.param.constant
     symbolic_value = transform.ParamConstantOp(
         transform.AnyParamType.get(), dummy_value
@@ -35,11 +35,11 @@ def testConstrainParamsOp(target):
         [symbolic_value], [smt.IntType.get()]
     )
     # CHECK-NEXT: ^bb{{.*}}(%[[PARAM_AS_SMT_SYMB:.*]]: !smt.int):
-    with InsertionPoint(constrain_params.body):
+    with ir.InsertionPoint(constrain_params.body):
         # CHECK: %[[C0:.*]] = smt.int.constant 0
-        c0 = smt.IntConstantOp(IntegerAttr.get(IntegerType.get_signless(32), 0))
+        c0 = smt.IntConstantOp(ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 0))
         # CHECK: %[[C43:.*]] = smt.int.constant 43
-        c43 = smt.IntConstantOp(IntegerAttr.get(IntegerType.get_signless(32), 43))
+        c43 = smt.IntConstantOp(ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 43))
         # CHECK: %[[LB:.*]] = smt.int.cmp le %[[C0]], %[[PARAM_AS_SMT_SYMB]]
         lb = smt.IntCmpOp(smt.IntPredicate.le, c0, constrain_params.body.arguments[0])
         # CHECK: %[[UB:.*]] = smt.int.cmp le %[[PARAM_AS_SMT_SYMB]], %[[C43]]