diff --git a/triton_viz/clients/sanitizer/sanitizer.py b/triton_viz/clients/sanitizer/sanitizer.py index 85291502..5c55464e 100644 --- a/triton_viz/clients/sanitizer/sanitizer.py +++ b/triton_viz/clients/sanitizer/sanitizer.py @@ -19,6 +19,7 @@ And, Or, Not, + Xor, sat, simplify, ) @@ -1085,6 +1086,9 @@ def eval(self) -> tuple[Union[ArithRef, list[ArithRef]], list[BoolRef]]: return expr, constraints def _to_z3(self) -> tuple[ArithRef, list]: + def _to_bool(x): + return x if isinstance(x, BoolRef) else (x != 0) + if self._z3 is not None: return self._z3, self._constraints @@ -1159,8 +1163,15 @@ def _to_z3(self) -> tuple[ArithRef, list]: self._z3 = If(lhs >= rhs, lhs, rhs) if self.op == "minimum": self._z3 = If(lhs <= rhs, lhs, rhs) - if self.op == "bitwise_and": - self._z3 = And(lhs, rhs) + if "bitwise" in self.op: + lhs = _to_bool(lhs) + rhs = _to_bool(rhs) + if self.op == "bitwise_and": + self._z3 = And(lhs, rhs) + if self.op == "bitwise_or": + self._z3 = Or(lhs, rhs) + if self.op == "bitwise_xor": + self._z3 = Xor(lhs, rhs) if self.op == "ashr": raise NotImplementedError( "Arithmetic shift right is not implemented in Z3"