kwonly replace arg

Author: makslevental

mlir/include/mlir-c/Bindings/Python/Interop.h

@@ -118,25 +118,24 @@
 
 /** Attribute on main C extension module (_mlir) that corresponds to the
  * type caster registration binding. The signature of the function is:
- *   def register_type_caster(MlirTypeID mlirTypeID, py::function typeCaster,
- *                              bool replace)
- * where replace indicates the typeCaster should replace any existing registered
- * type casters (such as those for upstream ConcreteTypes). The interface of the
- * typeCaster is:
- *   def type_caster(ir.Type) -> SubClassTypeT
- * where SubClassTypeT indicates the result should be a subclass (inherit from)
- * ir.Type.
+ *   def register_type_caster(MlirTypeID mlirTypeID, *, bool replace)
+ * which then takes a typeCaster (register_type_caster is meant to be used as a
+ * decorator from python), and where replace indicates the typeCaster should
+ * replace any existing registered type casters (such as those for upstream
+ * ConcreteTypes). The interface of the typeCaster is: def type_caster(ir.Type)
+ * -> SubClassTypeT where SubClassTypeT indicates the result should be a
+ * subclass (inherit from) ir.Type.
  */
 #define MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR "register_type_caster"
 
 /** Attribute on main C extension module (_mlir) that corresponds to the
  * value caster registration binding. The signature of the function is:
- *   def register_value_caster(MlirTypeID mlirTypeID, bool replace,
- *                              py::function valueCaster)
- * where replace indicates the valueCaster should replace any existing
- * registered value casters. The interface of the valueCaster is:
- *   def value_caster(ir.Value) -> SubClassValueT
- * where SubClassValueT indicates the result should be a subclass (inherit from)
+ *   def register_value_caster(MlirTypeID mlirTypeID, *, bool replace)
+ * which then takes a valueCaster (register_value_caster is meant to be used as
+ * a decorator, from python), and where replace indicates the valueCaster should
+ * replace any existing registered value casters. The interface of the
+ * valueCaster is: def value_caster(ir.Value) -> SubClassValueT where
+ * SubClassValueT indicates the result should be a subclass (inherit from)
  * ir.Value.
  */
 #define MLIR_PYTHON_CAPI_VALUE_CASTER_REGISTER_ATTR "register_value_caster"

mlir/include/mlir/Bindings/Python/PybindAdaptors.h

@@ -497,11 +497,10 @@ class mlir_type_subclass : public pure_subclass {
     if (getTypeIDFunction) {
       py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
           .attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)(
-              getTypeIDFunction(),
-              pybind11::cpp_function(
-                  [thisClass = thisClass](const py::object &mlirType) {
-                    return thisClass(mlirType);
-                  }));
+              getTypeIDFunction())(pybind11::cpp_function(
+              [thisClass = thisClass](const py::object &mlirType) {
+                return thisClass(mlirType);
+              }));
     }
   }
 };

mlir/lib/Bindings/Python/MainModule.cpp

@@ -44,7 +44,8 @@ PYBIND11_MODULE(_mlir, m) {
            "dialect_namespace"_a, "dialect_class"_a,
            "Testing hook for directly registering a dialect")
       .def("_register_operation_impl", &PyGlobals::registerOperationImpl,
-           "operation_name"_a, "operation_class"_a, "replace"_a = false,
+           "operation_name"_a, "operation_class"_a, py::kw_only(),
+           "replace"_a = false,
            "Testing hook for directly registering an operation");
 
   // Aside from making the globals accessible to python, having python manage
@@ -80,16 +81,19 @@ PYBIND11_MODULE(_mlir, m) {
               return opClass;
             });
       },
-      "dialect_class"_a, "replace"_a = false,
+      "dialect_class"_a, py::kw_only(), "replace"_a = false,
       "Produce a class decorator for registering an Operation class as part of "
       "a dialect");
   m.def(
       MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR,
-      [](MlirTypeID mlirTypeID, py::function typeCaster, bool replace) {
-        PyGlobals::get().registerTypeCaster(mlirTypeID, std::move(typeCaster),
-                                            replace);
+      [](MlirTypeID mlirTypeID, bool replace) -> py::cpp_function {
+        return py::cpp_function([mlirTypeID,
+                                 replace](py::object typeCaster) -> py::object {
+          PyGlobals::get().registerTypeCaster(mlirTypeID, typeCaster, replace);
+          return typeCaster;
+        });
       },
-      "typeid"_a, "type_caster"_a, "replace"_a = false,
+      "typeid"_a, py::kw_only(), "replace"_a = false,
       "Register a type caster for casting MLIR types to custom user types.");
   m.def(
       MLIR_PYTHON_CAPI_VALUE_CASTER_REGISTER_ATTR,
@@ -101,7 +105,7 @@ PYBIND11_MODULE(_mlir, m) {
               return valueCaster;
             });
       },
-      "typeid"_a, "replace"_a = false,
+      "typeid"_a, py::kw_only(), "replace"_a = false,
       "Register a value caster for casting MLIR values to custom user values.");
 
   // Define and populate IR submodule.

mlir/test/python/dialects/arith_dialect.py

@@ -41,9 +41,11 @@ def testFastMathFlags():
 def testArithValue():
     def _binary_op(lhs, rhs, op: str) -> "ArithValue":
         op = op.capitalize()
-        if arith._is_float_type(lhs.type):
+        if arith._is_float_type(lhs.type) and arith._is_float_type(rhs.type):
             op += "F"
-        elif arith._is_integer_like_type(lhs.type):
+        elif arith._is_integer_like_type(lhs.type) and arith._is_integer_like_type(
+            lhs.type
+        ):
             op += "I"
         else:
             raise NotImplementedError(f"Unsupported '{op}' operands: {lhs}, {rhs}")

mlir/test/python/dialects/python_test.py

@@ -508,31 +508,29 @@ def testCustomTypeTypeCaster():
         # CHECK: Type caster is already registered
         try:
 
+            @register_type_caster(c.typeid)
             def type_caster(pytype):
                 return test.TestIntegerRankedTensorType(pytype)
 
-            register_type_caster(c.typeid, type_caster)
         except RuntimeError as e:
             print(e)
 
-        def type_caster(pytype):
-            return RankedTensorType(pytype)
-
         # python_test dialect registers a caster for RankedTensorType in its extension (pybind) module.
         # So this one replaces that one (successfully). And then just to be sure we restore the original caster below.
-        register_type_caster(c.typeid, type_caster, replace=True)
+        @register_type_caster(c.typeid, replace=True)
+        def type_caster(pytype):
+            return RankedTensorType(pytype)
 
         d = tensor.EmptyOp([10, 10], IntegerType.get_signless(5)).result
         # CHECK: tensor<10x10xi5>
         print(d.type)
         # CHECK: ranked tensor type RankedTensorType(tensor<10x10xi5>)
         print("ranked tensor type", repr(d.type))
 
+        @register_type_caster(c.typeid, replace=True)
         def type_caster(pytype):
             return test.TestIntegerRankedTensorType(pytype)
 
-        register_type_caster(c.typeid, type_caster, replace=True)
-
         d = tensor.EmptyOp([10, 10], IntegerType.get_signless(5)).result
         # CHECK: tensor<10x10xi5>
         print(d.type)

mlir/test/python/lib/PythonTestModule.cpp

@@ -74,12 +74,11 @@ PYBIND11_MODULE(_mlirPythonTest, m) {
   MlirTypeID mlirRankedTensorTypeID = mlirRankedTensorTypeGetTypeID();
 
   py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
-      .attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)(
-          mlirRankedTensorTypeID,
+      .attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)(mlirRankedTensorTypeID,
+                                                        "replace"_a = true)(
           pybind11::cpp_function([typeCls](const py::object &mlirType) {
             return typeCls.get_class()(mlirType);
-          }),
-          /*replace=*/true);
+          }));
 
   auto valueCls = mlir_value_subclass(m, "TestTensorValue",
                                       mlirTypeIsAPythonTestTestTensorValue)