[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