From a06807d8dc6938eff3ba26db5dcec13c3379861a Mon Sep 17 00:00:00 2001 From: Vinicius Silva Date: Thu, 24 Sep 2026 01:02:03 +0000 Subject: [PATCH 1/4] arith: replace string predicates with CmpIPredicate enums in ArithValue comparisons --- .../mlir/extras/dialects/arith.py | 103 +++++++++++++----- 1 file changed, 73 insertions(+), 30 deletions(-) diff --git a/projects/eudsl-python-extras/mlir/extras/dialects/arith.py b/projects/eudsl-python-extras/mlir/extras/dialects/arith.py index a78e91965..b774af966 100644 --- a/projects/eudsl-python-extras/mlir/extras/dialects/arith.py +++ b/projects/eudsl-python-extras/mlir/extras/dialects/arith.py @@ -133,7 +133,13 @@ def index_cast( @register_attribute_builder("Arith_CmpIPredicateAttr", replace=True) -def _arith_CmpIPredicateAttr(predicate: Union[str, Attribute], context: Context): +def _arith_CmpIPredicateAttr( + predicate: Union[str, int, "CmpIPredicate", Attribute], context: Context +): + if isinstance(predicate, Attribute): + return predicate + if isinstance(predicate, CmpIPredicate): + return _arith_cmpipredicateattr(predicate, context) predicates = { "eq": CmpIPredicate.eq, "ne": CmpIPredicate.ne, @@ -156,15 +162,19 @@ def _arith_CmpIPredicateAttr(predicate: Union[str, Attribute], context: Context) 8: CmpIPredicate.ugt, 9: CmpIPredicate.uge, } - if isinstance(predicate, Attribute): - return predicate predicate = str(predicate) assert predicate in predicates, f"{predicate=} not in predicates" return _arith_cmpipredicateattr(predicates[predicate], context) @register_attribute_builder("Arith_CmpFPredicateAttr", replace=True) -def _arith_CmpFPredicateAttr(predicate: Union[str, Attribute], context: Context): +def _arith_CmpFPredicateAttr( + predicate: Union[str, int, "CmpFPredicate", Attribute], context: Context +): + if isinstance(predicate, Attribute): + return predicate + if isinstance(predicate, CmpFPredicate): + return _arith_cmpfpredicateattr(predicate, context) predicates = { "false": CmpFPredicate.AlwaysFalse, # ordered comparison @@ -190,13 +200,46 @@ def _arith_CmpFPredicateAttr(predicate: Union[str, Attribute], context: Context) # return always true "true": CmpFPredicate.AlwaysTrue, } - if isinstance(predicate, Attribute): - return predicate predicate = str(predicate) assert predicate in predicates, f"{predicate=} not in predicates" return _arith_cmpfpredicateattr(predicates[predicate], context) +# Canonical predicate keys are the *signed* CmpIPredicate enum members. +# Each ArithValue comparison dunder now passes one of these enum members +# (instead of an ad-hoc string like "lt") as `predicate`; _binary_op below +# resolves it to the correct concrete int/float, signed/unsigned enum. +_CMPI_PREDICATES = { + CmpIPredicate.eq: (CmpIPredicate.eq, CmpIPredicate.eq), + CmpIPredicate.ne: (CmpIPredicate.ne, CmpIPredicate.ne), + CmpIPredicate.slt: (CmpIPredicate.slt, CmpIPredicate.ult), + CmpIPredicate.sle: (CmpIPredicate.sle, CmpIPredicate.ule), + CmpIPredicate.sgt: (CmpIPredicate.sgt, CmpIPredicate.ugt), + CmpIPredicate.sge: (CmpIPredicate.sge, CmpIPredicate.uge), +} + +_CMPF_PREDICATES = { + CmpIPredicate.eq: (CmpFPredicate.OEQ, CmpFPredicate.UEQ), + CmpIPredicate.ne: (CmpFPredicate.ONE, CmpFPredicate.UNE), + CmpIPredicate.slt: (CmpFPredicate.OLT, CmpFPredicate.ULT), + CmpIPredicate.sle: (CmpFPredicate.OLE, CmpFPredicate.ULE), + CmpIPredicate.sgt: (CmpFPredicate.OGT, CmpFPredicate.UGT), + CmpIPredicate.sge: (CmpFPredicate.OGE, CmpFPredicate.UGE), +} + +# Maps a canonical predicate enum to the corresponding `operator` module +# attribute name, used only on the constant-folding path where we need to +# call e.g. operator.lt(lhs_literal, rhs_literal). +_CMP_PREDICATE_TO_OPERATOR_NAME = { + CmpIPredicate.eq: "eq", + CmpIPredicate.ne: "ne", + CmpIPredicate.slt: "lt", + CmpIPredicate.sle: "le", + CmpIPredicate.sgt: "gt", + CmpIPredicate.sge: "ge", +} + + def _binary_op( lhs: "ArithValue", rhs: "ArithValue", @@ -240,7 +283,7 @@ def _binary_op( # the corresponding operation on the literal values; e.g., operator.add. # note this is the same as op = operator.__dict__[op]. if predicate is not None: - op = predicate + op = _CMP_PREDICATE_TO_OPERATOR_NAME[predicate] op = operator.attrgetter(op)(operator) return klass(op(lhs, rhs), fold=True) @@ -272,23 +315,23 @@ def _binary_op( if predicate is not None: if isinstance(lhs.dtype, FloatType): - # ordered comparison - see above - predicate = "o" + predicate + enum_predicate = _CMPF_PREDICATES[predicate][0] else: assert isinstance( lhs.dtype, (IntegerType, IndexType) ), f"unsupported dtype for comparison: {lhs.dtype}" + signed_enum, unsigned_enum = _CMPI_PREDICATES[predicate] # eq, ne signs don't matter - if predicate not in {"eq", "ne"}: - if signedness is not None: - predicate = signedness + predicate - else: - if isinstance(lhs.dtype, IndexType) or lhs.dtype.is_unsigned: - predicate = "u" + predicate - else: - assert lhs.dtype.is_signed or lhs.dtype.is_signless - predicate = "s" + predicate - return lhs.__class__(op(predicate, lhs, rhs, loc=loc), dtype=lhs.dtype) + if predicate in {CmpIPredicate.eq, CmpIPredicate.ne}: + enum_predicate = signed_enum + elif signedness is not None: + enum_predicate = unsigned_enum if signedness == "u" else signed_enum + elif isinstance(lhs.dtype, IndexType) or lhs.dtype.is_unsigned: + enum_predicate = unsigned_enum + else: + assert lhs.dtype.is_signed or lhs.dtype.is_signless + enum_predicate = signed_enum + return lhs.__class__(op(enum_predicate, lhs, rhs, loc=loc), dtype=lhs.dtype) else: return lhs.__class__(op(lhs, rhs, loc=loc), dtype=lhs.dtype) @@ -394,7 +437,7 @@ def __eq__(self, other): return False if self is other: return True - return _binary_op(self, other, op="cmp", predicate="eq") + return _binary_op(self, other, op="cmp", predicate=CmpIPredicate.eq) def __ne__(self, other): if not isinstance(other, self.__class__): @@ -405,17 +448,17 @@ def __ne__(self, other): return True if self is other: return False - return _binary_op(self, other, op="cmp", predicate="ne") + return _binary_op(self, other, op="cmp", predicate=CmpIPredicate.ne) - __le__ = partialmethod(_binary_op, op="cmp", predicate="le") - __lt__ = partialmethod(_binary_op, op="cmp", predicate="lt") - __ge__ = partialmethod(_binary_op, op="cmp", predicate="ge") - __gt__ = partialmethod(_binary_op, op="cmp", predicate="gt") + __le__ = partialmethod(_binary_op, op="cmp", predicate=CmpIPredicate.sle) + __lt__ = partialmethod(_binary_op, op="cmp", predicate=CmpIPredicate.slt) + __ge__ = partialmethod(_binary_op, op="cmp", predicate=CmpIPredicate.sge) + __gt__ = partialmethod(_binary_op, op="cmp", predicate=CmpIPredicate.sgt) - __rle__ = partialmethod(_rbinary_op, op="cmp", predicate="le") - __rlt__ = partialmethod(_rbinary_op, op="cmp", predicate="lt") - __rge__ = partialmethod(_rbinary_op, op="cmp", predicate="ge") - __rgt__ = partialmethod(_rbinary_op, op="cmp", predicate="gt") + __rle__ = partialmethod(_rbinary_op, op="cmp", predicate=CmpIPredicate.sle) + __rlt__ = partialmethod(_rbinary_op, op="cmp", predicate=CmpIPredicate.slt) + __rge__ = partialmethod(_rbinary_op, op="cmp", predicate=CmpIPredicate.sge) + __rgt__ = partialmethod(_rbinary_op, op="cmp", predicate=CmpIPredicate.sgt) def _eq(self, other): return Value(self) == Value(other) @@ -537,4 +580,4 @@ class ArithCanonicalizer(Canonicalizer): function_patchers = [ArithPatchFunction] -canonicalizer = ArithCanonicalizer() +canonicalizer = ArithCanonicalizer() \ No newline at end of file From 37c82888e2cb5bedde79852319fe4ac5298043c3 Mon Sep 17 00:00:00 2001 From: Vinicius Silva Date: Fri, 25 Sep 2026 09:26:17 -0300 Subject: [PATCH 2/4] Update projects/eudsl-python-extras/mlir/extras/dialects/arith.py Co-authored-by: Maksim Levental --- projects/eudsl-python-extras/mlir/extras/dialects/arith.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/projects/eudsl-python-extras/mlir/extras/dialects/arith.py b/projects/eudsl-python-extras/mlir/extras/dialects/arith.py index b774af966..9612dca18 100644 --- a/projects/eudsl-python-extras/mlir/extras/dialects/arith.py +++ b/projects/eudsl-python-extras/mlir/extras/dialects/arith.py @@ -315,7 +315,7 @@ def _binary_op( if predicate is not None: if isinstance(lhs.dtype, FloatType): - enum_predicate = _CMPF_PREDICATES[predicate][0] + enum_predicate, _ = _CMPF_PREDICATES[predicate] else: assert isinstance( lhs.dtype, (IntegerType, IndexType) From e530e2824a6888e8cc01a1ea5069166fd6f883af Mon Sep 17 00:00:00 2001 From: Vinicius Silva Date: Fri, 25 Sep 2026 09:44:20 -0300 Subject: [PATCH 3/4] Remove blank line in ArithCanonicalizer class Removed an unnecessary blank line in the ArithCanonicalizer class. --- projects/eudsl-python-extras/mlir/extras/dialects/arith.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/projects/eudsl-python-extras/mlir/extras/dialects/arith.py b/projects/eudsl-python-extras/mlir/extras/dialects/arith.py index 9612dca18..9175be6df 100644 --- a/projects/eudsl-python-extras/mlir/extras/dialects/arith.py +++ b/projects/eudsl-python-extras/mlir/extras/dialects/arith.py @@ -579,5 +579,4 @@ class ArithCanonicalizer(Canonicalizer): cst_transformers = [CanonicalizeFMA] function_patchers = [ArithPatchFunction] - -canonicalizer = ArithCanonicalizer() \ No newline at end of file +canonicalizer = ArithCanonicalizer() From cdf6604e173f427a7ae8103cf3ad25da08b26233 Mon Sep 17 00:00:00 2001 From: Vinicius Silva Date: Sat, 26 Sep 2026 10:40:52 -0300 Subject: [PATCH 4/4] Instantiate ArithCanonicalizer in arith.py, removing a blank new Ensure proper instantiation of ArithCanonicalizer.