[Mlir-commits] [mlir] [mlir][python] fix `replace=True` for `register_operation` and `register_type_caster` (PR #70264)

Maksim Levental llvmlistbot at llvm.org
Wed Oct 25 17:34:24 PDT 2023


================
@@ -668,6 +671,27 @@ def testKnownOpView():
         # CHECK: OpView object
         print(repr(custom))
 
+        # constant should map to an extension OpView class in the arithmetic dialect.
+        constant = module.body.operations[3]
+        # CHECK: <mlir.dialects.arith.ConstantOp object
+        print(repr(constant))
+        # CHECK: literal value 0
+        print("literal value", constant.literal_value)
+
+        @_cext.register_operation(arith._Dialect, replace=True)
+        class ConstantOp(arith.ConstantOp):
+            def __init__(self, result, value, *, loc=None, ip=None):
+                if isinstance(value, int):
+                    super().__init__(IntegerAttr.get(result, value), loc=loc, ip=ip)
+                elif isinstance(value, float):
+                    super().__init__(FloatAttr.get(result, value), loc=loc, ip=ip)
+                else:
+                    super().__init__(value, loc=loc, ip=ip)
+
+        constant = module.body.operations[3]
+        # CHECK: <__main__.testKnownOpView.<locals>.ConstantOp object
+        print(repr(constant))
----------------
makslevental wrote:

"late" registration/replacement working correctly.

https://github.com/llvm/llvm-project/pull/70264


More information about the Mlir-commits mailing list