[Comb] Update Python binding decorators to wrap in the right order. (#…

…6377)

With a recent change upstream to the new `register_operation`
decorators, we accidentally broke the Python class hierarchy for comb
ops. Before this change, we would wrap the class generated by
`register_operation`. This meant the added `create` methods were
present on those classes, but when an operation was returned via
Python bindings methods, it would not be a subclass of our final
wrapped class; it would be the class returned by `register_operation`,
and `isinstance` would not work as expected.

With this change, we first wrap the class with our decorator, then
wrap it with `register_operation`. This way, when operations are
returned from the Python bindings, the returned operation is a
subclass of our final wrapped class, and `isinstance` works as
expected.

A somewhat hacky but thorough test is added that demonstrates this
issue, and confirms the new decorator order is working as expected.

integration_test/Bindings/Python/dialects/comb.py

@@ -176,3 +176,19 @@ def build(module):
     hw.HWModuleOp(name="test", body_builder=build)
 
   print(m)
+
+  # Check that isinstance is working properly for these ops, which broke with changes to op extensions.
+  for op in m.body.operations[0].body.blocks[0].operations:
+    # Only check comb, which has unusual class hierarchies.
+    if "comb." not in op.name:
+      continue
+
+    # Clean up the IR operation name to match the Python class name.
+    op_name = f"{op.name.split('.')[1].capitalize()}Op"
+    op_name = op_name.replace("cmp", "Cmp")
+    op_name = op_name.replace("sOp", "SOp")
+    op_name = op_name.replace("uOp", "UOp")
+
+    # Dynamically get the Python class, and check that the op isinstance of the class.
+    cls = getattr(comb, op_name)
+    assert isinstance(op, cls)

lib/Bindings/Python/dialects/comb.py

@@ -192,100 +192,100 @@ def create(low_bit, result_type, input=None):
                             needs_result_type=True)
 
 
-@UnaryOp
 @_ods_cext.register_operation(_Dialect, replace=True)
+@UnaryOp
 class ParityOp(ParityOp):
   pass
 
 
 # Sugar classes for the various non-variadic binary ops.
-@BinaryOp
 @_ods_cext.register_operation(_Dialect, replace=True)
+@BinaryOp
 class DivSOp(DivSOp):
   pass
 
 
-@BinaryOp
 @_ods_cext.register_operation(_Dialect, replace=True)
+@BinaryOp
 class DivUOp(DivUOp):
   pass
 
 
-@BinaryOp
 @_ods_cext.register_operation(_Dialect, replace=True)
+@BinaryOp
 class ModSOp(ModSOp):
   pass
 
 
-@BinaryOp
 @_ods_cext.register_operation(_Dialect, replace=True)
+@BinaryOp
 class ModUOp(ModUOp):
   pass
 
 
-@BinaryOp
 @_ods_cext.register_operation(_Dialect, replace=True)
+@BinaryOp
 class ShlOp(ShlOp):
   pass
 
 
-@BinaryOp
 @_ods_cext.register_operation(_Dialect, replace=True)
+@BinaryOp
 class ShrSOp(ShrSOp):
   pass
 
 
-@BinaryOp
 @_ods_cext.register_operation(_Dialect, replace=True)
+@BinaryOp
 class ShrUOp(ShrUOp):
   pass
 
 
-@BinaryOp
 @_ods_cext.register_operation(_Dialect, replace=True)
+@BinaryOp
 class SubOp(SubOp):
   pass
 
 
 # Sugar classes for the variadic ops.
-@VariadicOp
 @_ods_cext.register_operation(_Dialect, replace=True)
+@VariadicOp
 class AddOp(AddOp):
   pass
 
 
-@VariadicOp
 @_ods_cext.register_operation(_Dialect, replace=True)
+@VariadicOp
 class MulOp(MulOp):
   pass
 
 
-@VariadicOp
 @_ods_cext.register_operation(_Dialect, replace=True)
+@VariadicOp
 class AndOp(AndOp):
   pass
 
 
-@VariadicOp
 @_ods_cext.register_operation(_Dialect, replace=True)
+@VariadicOp
 class OrOp(OrOp):
   pass
 
 
-@VariadicOp
 @_ods_cext.register_operation(_Dialect, replace=True)
+@VariadicOp
 class XorOp(XorOp):
   pass
 
 
-@VariadicOp
 @_ods_cext.register_operation(_Dialect, replace=True)
+@VariadicOp
 class ConcatOp(ConcatOp):
   pass
 
 
 # Sugar classes for miscellaneous ops.
-@CreatableOp
 @_ods_cext.register_operation(_Dialect, replace=True)
+@CreatableOp
 class MuxOp(MuxOp):
   pass