[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