diff --git a/docs/architecture.md b/docs/architecture.md index 3853824..029ceba 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -33,7 +33,9 @@ flowchart LR - **optimize** (`optimize/`) — IR-level constant folding, dead-code elimination, and common-subexpression elimination on the scalar program. - **codegen** (`codegen/`) — emits readable Python source from the lowered - IR (SSA names, real loop structure). + IR (SSA names, real loop structure). Conditions proven to depend only + on loop indices select whole branches in Python; array conditions use + `xp.where`. - **exec + wrap** (`emit.py`) — compiles the source and wraps it so the namespace is taken from the call arguments (`array_api_compat.array_namespace`); `vec.source` and @@ -57,7 +59,9 @@ generated code calls small runtime helpers injected into its globals: for boolean/maybe-boolean/uint64 operands: implements the layered [exactness lattice](semantics.md#integer-exactness-lattice), including per-lane-exact uint64 `mod`/`floordiv` formulas for negative divisors - and result-aware subtraction. + and result-aware subtraction. Integer true division uses binary long + division with guard/sticky bits for correctly rounded float64 ratios, + without extracting array values into Python. The helpers live in `array_vectorize/runtime/` (`minmax.py`, `arith.py`, `dtype.py`), a stdlib-only leaf (enforced, with the other import diff --git a/docs/semantics.md b/docs/semantics.md index 80d2302..96c793a 100644 --- a/docs/semantics.md +++ b/docs/semantics.md @@ -15,6 +15,11 @@ enclosing where-select) and out of domain — live lanes are never touched disappear. Loop bodies and cross-function (helper) dead lanes are not tracked (conservative). +Conditions built only from constant-trip loop indices and literals select +one whole-array branch using Python selection. They do not promote the +selected branch to the unused branch's dtype. Array-dependent conditions +continue to select per lane with `xp.where`. + ### `and`/`or`/`not` are exact Numeric operands lower to value-selects @@ -40,6 +45,15 @@ call and checked at runtime on the actual values: 4. **float64** only when per-lane results fit no single dtype (a huge uint64 plus a negative signed value, or mixed-magnitude differences). +True division converts integer and boolean operands to floating arrays, +including inputs whose dtype is known only at call time. Integer ratios +round once to float64, without first rounding large integer operands. +This uses a fixed bit loop, which costs more at call time and first JIT +compilation than floating-point division. Negative integer +literals (also when assigned to local names) use the exact unsigned +remainder and floor-division paths. Unknown loop-carried entry values may +be boolean, so their arithmetic also uses runtime dtype dispatch. + ### Reject over miscompile Anything not provably translatable raises `VectorizationError` with all @@ -63,4 +77,4 @@ values. These are deliberate, tested, and inherent: | bool arithmetic is integer (`True + True == 2`) | exact: bool operands are cast to float64 (strict backends reject bool arithmetic) | | scalar exceptions inside `verify=` | expected `NaN` lanes | | mixed int/float `min`/`max` siblings | strict backends require matching array dtypes; literals adopt a provably-int sibling's dtype, else float64 | -| uint64 + signed arrays with mixed-sign per-lane results | float64 approximation (no single dtype holds both operand ranges) | +| integer arithmetic with results outside every integer array dtype | float64 approximation (including uint64 divided by a negative literal below int64-min) | diff --git a/src/array_vectorize/codegen/emitter.py b/src/array_vectorize/codegen/emitter.py index e4dcbd2..7bb29b5 100644 --- a/src/array_vectorize/codegen/emitter.py +++ b/src/array_vectorize/codegen/emitter.py @@ -79,7 +79,30 @@ def _gen_literal(lit: ir.Literal, ns: str) -> ast.expr: return ast.Constant(value=value) -def _gen_expr(node: ir.Node, ns: str) -> ast.expr: +def _is_python_scalar(node: ir.Node, names: frozenset[str]) -> bool: + """Prove raw scalar expressions built only from literals and loop indices.""" + match node: + case ir.Literal(): + return True + case ir.Ref(): + return node.name in names + case ir.BinOp(): + return ( + node.op in _BINOP_AST + and _is_python_scalar(node.left, names) + and _is_python_scalar(node.right, names) + ) + case ir.UnaryOp(): + return _is_python_scalar(node.operand, names) + case ir.Compare(): + return _is_python_scalar(node.left, names) and _is_python_scalar(node.right, names) + case ir.Logical(): + return all(_is_python_scalar(part, names) for part in node.parts) + case _: + return False + + +def _gen_expr(node: ir.Node, ns: str, scalar_names: frozenset[str] = frozenset()) -> ast.expr: match node: case ir.Literal(): return _gen_literal(node, ns) @@ -91,43 +114,69 @@ def _gen_expr(node: ir.Node, ns: str) -> ast.expr: if node.op in _BINOP_AST: return ast.BinOp( op=_BINOP_AST[node.op](), - left=_gen_expr(node.left, ns), - right=_gen_expr(node.right, ns), + left=_gen_expr(node.left, ns, scalar_names), + right=_gen_expr(node.right, ns, scalar_names), ) return _xp_call( - ns, _BINOP_XP[node.op], [_gen_expr(node.left, ns), _gen_expr(node.right, ns)] + ns, + _BINOP_XP[node.op], + [_gen_expr(node.left, ns, scalar_names), _gen_expr(node.right, ns, scalar_names)], ) case ir.UnaryOp(): if node.op == "neg": - return ast.UnaryOp(op=ast.USub(), operand=_gen_expr(node.operand, ns)) + return ast.UnaryOp(op=ast.USub(), operand=_gen_expr(node.operand, ns, scalar_names)) if node.op == "pos": - return ast.UnaryOp(op=ast.UAdd(), operand=_gen_expr(node.operand, ns)) + return ast.UnaryOp(op=ast.UAdd(), operand=_gen_expr(node.operand, ns, scalar_names)) if node.op == "invert": - return ast.UnaryOp(op=ast.Invert(), operand=_gen_expr(node.operand, ns)) - return _xp_call(ns, "logical_not", [_gen_expr(node.operand, ns)]) + return ast.UnaryOp( + op=ast.Invert(), operand=_gen_expr(node.operand, ns, scalar_names) + ) + if _is_python_scalar(node.operand, scalar_names): + return ast.UnaryOp(op=ast.Not(), operand=_gen_expr(node.operand, ns, scalar_names)) + return _xp_call(ns, "logical_not", [_gen_expr(node.operand, ns, scalar_names)]) case ir.Compare(): return ast.Compare( - left=_gen_expr(node.left, ns), + left=_gen_expr(node.left, ns, scalar_names), ops=[_CMPOP_AST[node.op]()], - comparators=[_gen_expr(node.right, ns)], + comparators=[_gen_expr(node.right, ns, scalar_names)], ) case ir.Logical(): + if _is_python_scalar(node, scalar_names): + return ast.BoolOp( + op=ast.And() if node.op == "and" else ast.Or(), + values=[_gen_expr(part, ns, scalar_names) for part in node.parts], + ) fn = "logical_and" if node.op == "and" else "logical_or" - folded = _gen_expr(node.parts[0], ns) + folded = _gen_expr(node.parts[0], ns, scalar_names) for part in node.parts[1:]: - folded = _xp_call(ns, fn, [folded, _gen_expr(part, ns)]) + folded = _xp_call(ns, fn, [folded, _gen_expr(part, ns, scalar_names)]) return folded case ir.Where(): + if _is_python_scalar(node.cond, scalar_names): + # A loop-index condition selects one whole-array branch. + # Python selection preserves its dtype without sending a + # raw bool or incompatible branch dtypes to xp.where (D6). + return ast.IfExp( + test=_gen_expr(node.cond, ns, scalar_names), + body=_gen_expr(node.then, ns, scalar_names), + orelse=_gen_expr(node.other, ns, scalar_names), + ) return _xp_call( ns, "where", - [_gen_expr(node.cond, ns), _gen_expr(node.then, ns), _gen_expr(node.other, ns)], + [ + _gen_expr(node.cond, ns, scalar_names), + _gen_expr(node.then, ns, scalar_names), + _gen_expr(node.other, ns, scalar_names), + ], ) case ir.Call(): - return _xp_call(ns, node.fn, [_gen_expr(a, ns) for a in node.args]) + return _xp_call(ns, node.fn, [_gen_expr(a, ns, scalar_names) for a in node.args]) case ir.FuncCall(): return ast.Call( - func=_load(node.fn), args=[_gen_expr(a, ns) for a in node.args], keywords=[] + func=_load(node.fn), + args=[_gen_expr(a, ns, scalar_names) for a in node.args], + keywords=[], ) case _: assert_never(node) @@ -200,15 +249,17 @@ def _gen_range_call(loop: ir.Loop, ns: str) -> ast.Call: return ast.Call(func=_load("range"), args=args, keywords=[]) -def _gen_stmt(stmt: ir.Stmt, ns: str) -> ast.stmt: +def _gen_stmt(stmt: ir.Stmt, ns: str, scalar_names: frozenset[str] = frozenset()) -> ast.stmt: match stmt: case ir.Binding(): return ast.Assign( targets=[ast.Name(id=stmt.name, ctx=ast.Store())], - value=_gen_expr(stmt.expr, ns), + value=_gen_expr(stmt.expr, ns, scalar_names), ) case ir.Loop(): - body: list[ast.stmt] = [_gen_stmt(inner, ns) for inner in stmt.body] + body: list[ast.stmt] = [ + _gen_stmt(inner, ns, scalar_names | {stmt.var}) for inner in stmt.body + ] if not body: # (15) a retained loop with a fully dead body still needs a statement body = [ast.Pass()] diff --git a/src/array_vectorize/lower/expressions.py b/src/array_vectorize/lower/expressions.py index e6a31ef..3d7549c 100644 --- a/src/array_vectorize/lower/expressions.py +++ b/src/array_vectorize/lower/expressions.py @@ -15,6 +15,7 @@ from array_vectorize import ir from array_vectorize.frontend import tables from array_vectorize.lower import sanitize +from array_vectorize.optimize import constfold from array_vectorize.runtime import registry __all__ = ["_ExpressionLowerer"] @@ -344,22 +345,46 @@ def _maybe_bool_result(self, node: ir.Node) -> bool: def _lower_binop(self, op: ast.operator, left: ir.Node, right: ir.Node) -> ir.Node: """Build a BinOp with the exactness fixes shared by `a op b` and `a op= b`.""" + if isinstance(op, ast.FloorDiv | ast.Mod): + literal_left, literal_right = self._literal_value(left), self._literal_value(right) + if literal_left is not None and literal_right is not None: + # Keep foldable constants usable as loop bounds, while + # zero divisors and overflow retain runtime dispatch. + folded = constfold._fold_binop(_BINOPS[type(op)], literal_left, literal_right) + if folded is not None: + return folded # Python bool arithmetic is integer (True + True == 2), but array # backends either saturate (bool + bool == True) or reject bools in - # arithmetic outright (strict backends). Cast bool-typed operands to - # float64: exact for bools (0/1) and ints up to 2**53, and a no-op - # for floats — so loop-carried variables whose kind changes across - # iterations stay correct too. + # arithmetic outright (strict backends). Runtime dispatch chooses + # a numeric dtype without rounding integer operands, including + # loop-carried values whose kind changes across iterations. if isinstance(op, ast.Add | ast.Sub | ast.Mult | ast.Div | ast.FloorDiv | ast.Mod) and ( "bool" in (self._numeric_kind(left), self._numeric_kind(right)) or self._maybe_bool_result(left) or self._maybe_bool_result(right) + # True division requires floating operands even when input + # kinds are unknown. Negative integer literals cannot be + # promoted into uint64 by a bare backend operator. + or ( + isinstance(op, ast.Div) + and (self._numeric_kind(left) != "float" or self._numeric_kind(right) != "float") + ) + or ( + isinstance(op, ast.FloorDiv | ast.Mod) + and ( + self._is_negative_int(self._literal_value(left) or left) + or self._is_negative_int(self._literal_value(right) or right) + ) + ) ): # the whole operation goes through the runtime _vec_arith # helper: bools need a numeric representation, strict # backends require single-dtype or float operands, and the # uint64 integer paths are exact per lane (modular for # add/sub/mul; int64 magnitudes for negative divisors) + if isinstance(op, ast.FloorDiv | ast.Mod): + left = self._literal_value(left) or left + right = self._literal_value(right) or right ir_op = _BINOPS[type(op)] return ir.FuncCall( self.arith_name, diff --git a/src/array_vectorize/lower/kinds.py b/src/array_vectorize/lower/kinds.py index cbacb94..518ee70 100644 --- a/src/array_vectorize/lower/kinds.py +++ b/src/array_vectorize/lower/kinds.py @@ -17,8 +17,8 @@ Python scalar (parameter default, loop variable, literal binding, or a local computed from one). Operators promote scalars fine; ``xp.*`` function arguments do not (strict backends reject plain scalars). -- ``maybe_bool`` (was ``_maybe_bool_names``): the value is a min/max - result with unknown static kind whose runtime dtype may be boolean, so +- ``maybe_bool`` (was ``_maybe_bool_names``): an unknown min/max result + or loop-carried entry value may have boolean dtype at runtime, so arithmetic must use the runtime-polymorphic intify. - carried-assign kinds (was ``_carried_assign_kinds``): per active loop, the kinds assigned to loop-carried names in the body; a carried diff --git a/src/array_vectorize/lower/loops.py b/src/array_vectorize/lower/loops.py index 64632fe..7a2ba4f 100644 --- a/src/array_vectorize/lower/loops.py +++ b/src/array_vectorize/lower/loops.py @@ -125,7 +125,9 @@ def lower_for(self, stmt: ast.For) -> None: # scalar stays a raw scalar, so track that for call sites if self.kinds.is_scalar(pre_name): self.kinds.mark_scalar(loop_carried_name) - if self.kinds.is_maybe_bool(pre_name): + if self.kinds.kind(pre_name) is None or self.kinds.is_maybe_bool(pre_name): + # Unknown entry values can be boolean arrays even when + # every body assignment has the same unknown static kind. self.kinds.mark_maybe_bool(loop_carried_name) phi_end = len(self.bindings) diff --git a/src/array_vectorize/lower/statements.py b/src/array_vectorize/lower/statements.py index aa86741..b186db1 100644 --- a/src/array_vectorize/lower/statements.py +++ b/src/array_vectorize/lower/statements.py @@ -46,8 +46,8 @@ def lower_assign(self, stmt: ast.Assign | ast.AugAssign) -> None: self.definite[var] = name kind = self._numeric_kind(value) self.kinds.set_kind(name, kind) - if isinstance(value, ir.Literal): - self.kinds.set_literal(name, value) + if (literal := self._literal_value(value)) is not None: + self.kinds.set_literal(name, literal) else: self.kinds.drop_literal(name) if self._possibly_scalar(value): diff --git a/src/array_vectorize/runtime/arith.py b/src/array_vectorize/runtime/arith.py index e89bb51..ccd2e94 100644 --- a/src/array_vectorize/runtime/arith.py +++ b/src/array_vectorize/runtime/arith.py @@ -107,6 +107,90 @@ def _u64_sub_exact(xp: Any, left: Any, right: Any) -> Any | None: return None +def _integer_exponent(xp: Any, value: Any) -> Any: + """Highest set bit of an unsigned magnitude stored in int64; zero maps to zero.""" + exponent = xp.zeros_like(value) + for shift in (32, 16, 8, 4, 2, 1): + higher = (value >> shift) & ((1 << (64 - shift)) - 1) + nonzero = higher != 0 + exponent = exponent + xp.astype(nonzero, xp.int64) * shift + value = xp.where(nonzero, higher, value) + return exponent + + +def _integer_true_divide(xp: Any, left: Any, right: Any) -> Any: + """Correctly round ratios of 64-bit integers without rounding inputs. + + Binary long division collects 53 significand bits plus guard/sticky + bits. Signed int64 stores unsigned magnitudes modularly, allowing + backends without uint64 arithmetic to use the same path. Remainders + are doubled via subtraction so comparisons never depend on overflow. + The fixed trip count also permits tracing by JAX and PyTorch. + """ + left_array = xp.asarray(left) + right_array = xp.asarray(right) + a = xp.astype(left_array, xp.int64) + b = xp.astype(right_array, xp.int64) + l_neg = a < 0 if "uint" not in str(left_array.dtype) else xp.asarray(False) + r_neg = b < 0 if "uint" not in str(right_array.dtype) else xp.asarray(False) + a = xp.where(l_neg, -a, a) + b = xp.where(r_neg, -b, b) + zero_divisor = b == 0 + b = xp.where(zero_divisor, xp.asarray(1, dtype=xp.int64), b) + leading_position = _integer_exponent(xp, a) - _integer_exponent(xp, b) + one = xp.asarray(1, dtype=xp.int64) + zero = xp.asarray(0, dtype=xp.int64) + # Skip quotient-leading zeros exactly. The first quotient bit lies + # at leading_position or one place below it. For a nonnegative + # position, preconsume numerator bits above it; otherwise preconsume + # all numerator bits and the leading fractional zeros. + shift = xp.where(leading_position < 0, zero, leading_position + 1) + shift = xp.where(shift > 63, xp.asarray(63, dtype=xp.int64), shift) + mask = ((one << (63 - shift)) - 1) * 2 + 1 + prefix = (a >> shift) & mask + prefix = xp.where(leading_position == 63, zero, prefix) + left_shift = xp.where(leading_position >= 0, zero, -leading_position - 1) + remainder = xp.where(leading_position >= 0, prefix, a << left_shift) + remainder = remainder + xp.zeros_like(a + b) + mantissa = xp.zeros_like(remainder) + count = xp.zeros_like(remainder) + exponent = xp.zeros_like(remainder) + guard = xp.zeros_like(remainder, dtype=xp.bool) + sticky = xp.zeros_like(guard) + # 53 significand bits, one guard bit, and at most one leading zero. + for offset in range(55): + position = leading_position - offset + safe_position = xp.where(position < 0, zero, position) + incoming = xp.where(position >= 0, (a >> safe_position) & 1, zero) + threshold = b - remainder - incoming + # Unsigned comparison using signed storage: a set sign bit means + # the unsigned value is larger than every nonnegative int64. + bit = xp.where((remainder < 0) != (threshold < 0), remainder < 0, remainder >= threshold) + remainder = xp.where( + bit, remainder - (b - remainder) + incoming, remainder + remainder + incoming + ) + started = (count != 0) | bit + exponent = xp.where((count == 0) & bit, position, exponent) + collect = started & (count < 53) + mantissa = xp.where(collect, mantissa * 2 + xp.astype(bit, xp.int64), mantissa) + rounding_bit = count == 53 + guard = xp.where(rounding_bit, bit, guard) + unconsumed = (a & ((one << safe_position) - 1)) != 0 + sticky = xp.where(rounding_bit, (remainder != 0) | unconsumed, sticky) + count = xp.where(started, count + 1, count) + round_up = guard & (sticky | ((mantissa & 1) != 0)) + mantissa = mantissa + xp.astype(round_up, xp.int64) + result = xp.astype(mantissa, xp.float64) * xp.pow( + xp.asarray(2.0, dtype=xp.float64), xp.astype(exponent - 52, xp.float64) + ) + result = xp.where(l_neg != r_neg, -result, result) + # Numeric exceptions deliberately become IEEE values (design D3). + denominator = xp.where( + zero_divisor, xp.asarray(0.0, dtype=xp.float64), xp.asarray(1.0, dtype=xp.float64) + ) + return result / denominator + + def _vec_arith(xp: Any, op_id: Any, left: Any, right: Any = None) -> Any: """Exact-semantics arithmetic for bool-intified operations. @@ -122,11 +206,9 @@ def _vec_arith(xp: Any, op_id: Any, left: Any, right: Any = None) -> Any: - add/sub/mul and floordiv/mod with non-negative divisors: modular uint64 (exact while the true result is in [0, 2**64) — all any dtype can hold); - - array % negative-literal and array // negative-literal: the results - ALWAYS fit int64 (remainder magnitude < |divisor| <= 2**63; floor - magnitude <= 2**63), so the magnitudes are computed exactly in - uint64 and negated into int64 — per lane, regardless of how large - the input values are; + - array % negative-literal and array // negative-literal: magnitudes + are computed exactly in uint64, then negated into int64 when they + fit. Quotients below int64-min use the float64 fallback; - negative-literal % array and negation: int64 whenever the values fit (checked at runtime); otherwise float64, because those results are genuinely unrepresentable in int64 on the offending lanes. @@ -137,6 +219,8 @@ def _vec_arith(xp: Any, op_id: Any, left: Any, right: Any = None) -> Any: has_real_float = any( hasattr(a, "dtype") and ("float" in str(a.dtype) or "complex" in str(a.dtype)) for a in args ) or any(isinstance(a, float) for a in args) + if op == "div" and not has_real_float: + return _integer_true_divide(xp, left, right) if any(_is_u64(a) for a in args) and not has_real_float: others_intish = all( (hasattr(a, "dtype") and ("uint" in str(a.dtype) or "bool" in str(a.dtype))) @@ -197,7 +281,13 @@ def _vec_arith(xp: Any, op_id: Any, left: Any, right: Any = None) -> Any: q = left // d r = left % d mag = q + xp.astype(r > 0, left.dtype) # ceil(a / |d|) - return -xp.astype(mag, xp.int64) + if bool(xp.all(mag <= 2**63)): + return -xp.astype(mag, xp.int64) + # Larger negative quotients cannot fit int64. Preserve + # their sign with the lattice's float64 fallback. + return xp.where( + mag == 0, xp.asarray(0.0, dtype=xp.float64), -xp.astype(mag, xp.float64) + ) if ( op == "mod" and isinstance(left, int) @@ -234,9 +324,6 @@ def _vec_arith(xp: Any, op_id: Any, left: Any, right: Any = None) -> Any: if bool(xp.all(left <= 2**63)): return -xp.astype(left, xp.int64) return -xp.astype(left, xp.float64) - if op == "div": - divide: Callable[[Any, Any], Any] = operator.truediv - return divide(xp.asarray(left, dtype=xp.float64), xp.asarray(right, dtype=xp.float64)) # generic: cast to the common dtype and compute classes = [dtype._describe_dtype(a) for a in args] has_float = any(f for f, _ in classes) diff --git a/tests/backends/test_array_api_strict.py b/tests/backends/test_array_api_strict.py index d2f4520..f81634a 100644 --- a/tests/backends/test_array_api_strict.py +++ b/tests/backends/test_array_api_strict.py @@ -432,3 +432,96 @@ def test_uint64_truediv_negative_divisor_strict() -> None: vec = array_vectorize.vectorize(support.make_fn(" a = max(x, True)\n return a / -2")) got = vec(xps.asarray([3], dtype=xps.uint64)) assert float(got[0]) == -1.5 + + +@pytest.mark.parametrize("op", ["%", "//"]) +@pytest.mark.parametrize("bound", [False, True]) +def test_negative_integer_with_unsigned_input(op: str, bound: bool) -> None: + import array_api_strict as xps + + body = f" return -7 {op} x" + if bound: + body = f" n = -7\n return n {op} x" + scalar = support.make_fn(body) + values = [1, 2, 2**63, 2**64 - 1] + got = array_vectorize.vectorize(scalar)(xps.asarray(values, dtype=xps.uint64)) + assert [int(v) for v in got] == [scalar(v) for v in values] + assert np.asarray(got).dtype.kind == ("u" if op == "%" else "i") + + +@pytest.mark.parametrize("dtype", ["bool", "int64", "uint64"]) +def test_true_division_of_integer_inputs(dtype: str) -> None: + import array_api_strict as xps + + scalar = support.make_fn(" return x / y", defaults="x, y") + values: list[int] = [True, False] if dtype == "bool" else [1, 2, 3] + if dtype == "uint64": + values += [2**63, 2**64 - 1] + denominator = True if dtype == "bool" else 3 + got = array_vectorize.vectorize(scalar)( + xps.asarray(values, dtype=getattr(xps, dtype)), + xps.asarray([denominator] * len(values), dtype=getattr(xps, dtype)), + ) + assert [float(v) for v in got] == [scalar(v, denominator) for v in values] + + +@pytest.mark.parametrize("assignment", ["y = y + 1", "y += 1", "y = y + y"]) +@pytest.mark.parametrize("trips", [0, 1, 3]) +def test_boolean_input_carried_through_loop(assignment: str, trips: int) -> None: + import array_api_strict as xps + + scalar = support.make_fn( + f" y = x\n for i in range({trips}):\n {assignment}\n return y" + ) + got = array_vectorize.vectorize(scalar)(xps.asarray([False, True])) + assert [int(v) for v in got] == [scalar(False), scalar(True)] + assert np.asarray(got).dtype.kind == ("b" if trips == 0 else "i") + + +@pytest.mark.parametrize( + "condition", ["i == 0", "i == 0 and i < 2", "i == 0 or i < 0", "not (i != 0)", "i + 1 == 1"] +) +@pytest.mark.parametrize("statement", [False, True]) +@pytest.mark.parametrize("trips", [1, 3]) +def test_loop_index_selects_branch_without_dtype_promotion( + condition: str, statement: bool, trips: int +) -> None: + import array_api_strict as xps + + body = f" y = x if {condition} else y + 1" + if statement: + body = f" if {condition}:\n y = x\n else:\n y = y + 1" + scalar = support.make_fn(f" y = True\n for i in range({trips}):\n{body}\n return y") + values = [-0.0, 0.0, -1.0, 1.0, np.nextafter(0.0, 1.0), -np.inf, np.inf, np.nan] + got = array_vectorize.vectorize(scalar)(xps.asarray(values, dtype=xps.float64)) + np.testing.assert_equal(np.asarray(got), [scalar(v) for v in values]) + assert np.asarray(got).dtype.kind == "f" + if trips == 1: + np.testing.assert_equal(np.signbit(np.asarray(got)), np.signbit(values)) + + +def test_integer_division_does_not_round_the_denominator_first() -> None: + import array_api_strict as xps + + scalar = support.make_fn(" return x / y", defaults="x, y") + left = [1, -1, -(2**63), 2**63 - 1, 0] + right = [2**53 + 1, 2**53 + 1, -1, 3, -1] + got = array_vectorize.vectorize(scalar)( + xps.asarray(left, dtype=xps.int64), + xps.asarray(right, dtype=xps.int64), + ) + expected = [a / b for a, b in zip(left, right, strict=True)] + np.testing.assert_array_equal(np.asarray(got), expected) + np.testing.assert_array_equal(np.signbit(np.asarray(got)), np.signbit(expected)) + + +def test_unsigned_floor_division_preserves_unrepresentable_negative_sign() -> None: + import array_api_strict as xps + + scalar = support.make_fn(" return x // -1") + values = [0, 1, 2**63, 2**63 + 1, 2**64 - 1] + got = array_vectorize.vectorize(scalar)(xps.asarray(values, dtype=xps.uint64)) + # No integer array dtype holds these results; the lattice uses float64. + assert np.asarray(got).dtype.kind == "f" + np.testing.assert_array_equal(np.asarray(got), [float(-v) for v in values]) + np.testing.assert_array_equal(np.signbit(np.asarray(got)), [v != 0 for v in values]) diff --git a/tests/backends/test_jax_torch_compilation.py b/tests/backends/test_jax_torch_compilation.py index d8eccfb..b8cab11 100644 --- a/tests/backends/test_jax_torch_compilation.py +++ b/tests/backends/test_jax_torch_compilation.py @@ -256,3 +256,29 @@ def test_closure_captured_jax_array() -> None: assert vec.__kwdefaults__["T"] is mod.T got = vec(mod.T * 10.0) assert np.allclose(np.asarray(got), [11.0, 22.0, 33.0]) + + +@pytest.mark.timeout(600) +@pytest.mark.parametrize("backend", ["jax", "torch"]) +def test_compiled_integer_division_preserves_python_rounding(backend: str) -> None: + scalar = support.make_fn(" return x / y", defaults="x, y") + left = [1, -1, -(2**63), 2**63 - 1, 0] + right = [2**53 + 1, 2**53 + 1, -1, 3, -1] + if backend == "jax": + jax = pytest.importorskip("jax") + jax.config.update("jax_enable_x64", True) + import jax.numpy as jnp + + vec = array_vectorize.vectorize(scalar, namespace=jnp) + got = jax.jit(vec)(jnp.asarray(left, dtype=jnp.int64), jnp.asarray(right, dtype=jnp.int64)) + else: + torch = pytest.importorskip("torch") + import array_api_compat.torch + + vec = array_vectorize.vectorize(scalar, namespace=array_api_compat.torch) + got = torch.compile(vec)( + torch.tensor(left, dtype=torch.int64), torch.tensor(right, dtype=torch.int64) + ) + expected = [a / b for a, b in zip(left, right, strict=True)] + np.testing.assert_array_equal(np.asarray(got), expected) + np.testing.assert_array_equal(np.signbit(np.asarray(got)), np.signbit(expected)) diff --git a/tests/behavior/test_arithmetic_dtypes.py b/tests/behavior/test_arithmetic_dtypes.py index d90376c..35d36f8 100644 --- a/tests/behavior/test_arithmetic_dtypes.py +++ b/tests/behavior/test_arithmetic_dtypes.py @@ -913,3 +913,53 @@ def test_nonnegative_signed_left_mod_uint64_exact() -> None: ) got = vec(np.asarray([2**63 + 1], dtype=np.uint64), np.asarray([7], dtype=np.int64)) assert int(got[0]) == 7 + + +@pytest.mark.parametrize("op", ["%", "//"]) +@pytest.mark.parametrize("bound", [False, True]) +def test_negative_integer_with_unsigned_input(op: str, bound: bool) -> None: + body = f" return -7 {op} x" + if bound: + body = f" n = -7\n return n {op} x" + scalar = support.make_fn(body) + values = [1, 2, 2**63, 2**64 - 1] + got = array_vectorize.vectorize(scalar)(np.asarray(values, dtype=np.uint64)) + assert [int(v) for v in got] == [scalar(v) for v in values] + assert np.asarray(got).dtype.kind == ("u" if op == "%" else "i") + + +@pytest.mark.parametrize("dtype", ["bool", "int64", "uint64"]) +def test_true_division_of_integer_inputs(dtype: str) -> None: + scalar = support.make_fn(" return x / y", defaults="x, y") + values: list[int] = [True, False] if dtype == "bool" else [1, 2, 3] + if dtype == "uint64": + values += [2**63, 2**64 - 1] + denominator = True if dtype == "bool" else 3 + got = array_vectorize.vectorize(scalar)( + np.asarray(values, dtype=getattr(np, dtype)), + np.asarray([denominator] * len(values), dtype=getattr(np, dtype)), + ) + assert [float(v) for v in got] == [scalar(v, denominator) for v in values] + + +def test_integer_division_does_not_round_the_denominator_first() -> None: + scalar = support.make_fn(" return x / y", defaults="x, y") + left = [1, -1, -(2**63), 2**63 - 1, 0] + right = [2**53 + 1, 2**53 + 1, -1, 3, -1] + got = array_vectorize.vectorize(scalar)( + np.asarray(left, dtype=np.int64), + np.asarray(right, dtype=np.int64), + ) + expected = [a / b for a, b in zip(left, right, strict=True)] + np.testing.assert_array_equal(np.asarray(got), expected) + np.testing.assert_array_equal(np.signbit(np.asarray(got)), np.signbit(expected)) + + +def test_unsigned_floor_division_preserves_unrepresentable_negative_sign() -> None: + scalar = support.make_fn(" return x // -1") + values = [0, 1, 2**63, 2**63 + 1, 2**64 - 1] + got = array_vectorize.vectorize(scalar)(np.asarray(values, dtype=np.uint64)) + # No integer array dtype holds these results; the lattice uses float64. + assert np.asarray(got).dtype.kind == "f" + np.testing.assert_array_equal(np.asarray(got), [float(-v) for v in values]) + np.testing.assert_array_equal(np.signbit(np.asarray(got)), [v != 0 for v in values]) diff --git a/tests/behavior/test_loops.py b/tests/behavior/test_loops.py index 2d32f50..261cb4c 100644 --- a/tests/behavior/test_loops.py +++ b/tests/behavior/test_loops.py @@ -458,3 +458,40 @@ def test_assign_loop_var_rejected() -> None: array_vectorize.vectorize( support.make_fn(" for i in range(3):\n i = i + 1\n return x") ) + + +@pytest.mark.parametrize("assignment", ["y = y + 1", "y += 1", "y = y + y"]) +@pytest.mark.parametrize("trips", [0, 1, 3]) +def test_boolean_input_carried_through_loop(assignment: str, trips: int) -> None: + scalar = support.make_fn( + f" y = x\n for i in range({trips}):\n {assignment}\n return y" + ) + got = array_vectorize.vectorize(scalar)(np.asarray([False, True])) + assert [int(v) for v in got] == [scalar(False), scalar(True)] + assert np.asarray(got).dtype.kind == ("b" if trips == 0 else "i") + + +@pytest.mark.parametrize( + "condition", ["i == 0", "i == 0 and i < 2", "i == 0 or i < 0", "not (i != 0)", "i + 1 == 1"] +) +@pytest.mark.parametrize("statement", [False, True]) +@pytest.mark.parametrize("trips", [1, 3]) +def test_loop_index_selects_branch_without_dtype_promotion( + condition: str, statement: bool, trips: int +) -> None: + body = f" y = x if {condition} else y + 1" + if statement: + body = f" if {condition}:\n y = x\n else:\n y = y + 1" + scalar = support.make_fn(f" y = True\n for i in range({trips}):\n{body}\n return y") + values = [-0.0, 0.0, -1.0, 1.0, np.nextafter(0.0, 1.0), -np.inf, np.inf, np.nan] + got = array_vectorize.vectorize(scalar)(np.asarray(values, dtype=np.float64)) + np.testing.assert_equal(np.asarray(got), [scalar(v) for v in values]) + assert np.asarray(got).dtype.kind == "f" + if trips == 1: + np.testing.assert_equal(np.signbit(np.asarray(got)), np.signbit(values)) + + +@pytest.mark.parametrize(("bounds", "trips"), [("-7 % 3", 2), ("7 // -3, 0", 3)]) +def test_negative_constant_arithmetic_in_loop_bounds(bounds: str, trips: int) -> None: + scalar = support.make_fn(f" for i in range({bounds}):\n x = x + 1\n return x") + np.testing.assert_array_equal(array_vectorize.vectorize(scalar)(np.asarray([0])), [trips]) diff --git a/tests/golden/cases/arith_ops.py b/tests/golden/cases/arith_ops.py index 0d5e39e..b43a6d1 100644 --- a/tests/golden/cases/arith_ops.py +++ b/tests/golden/cases/arith_ops.py @@ -11,4 +11,4 @@ def arith_ops(x, y): return (x + y - x * y) / (x * x + y * y) """ xp = array_namespace(x, y) - return (x + y - x * y) / (x * x + y * y) + return _vec_arith(xp, 4, x + y - x * y, x * x + y * y) diff --git a/tests/golden/cases/complex_expr.py b/tests/golden/cases/complex_expr.py index bb895ff..e5771b1 100644 --- a/tests/golden/cases/complex_expr.py +++ b/tests/golden/cases/complex_expr.py @@ -13,4 +13,4 @@ def complex_expr(x, y): """ xp = array_namespace(x, y) t = xp.exp(xp.astype(xp.asarray(-(x * x + y * y)), xp.float64)) - return t / (1.0 + t) + return _vec_arith(xp, 4, t, 1.0 + t) diff --git a/tests/golden/cases/ternary.py b/tests/golden/cases/ternary.py index 6c7826f..24cc3fb 100644 --- a/tests/golden/cases/ternary.py +++ b/tests/golden/cases/ternary.py @@ -11,4 +11,4 @@ def ternary(x): return x * 2 if x > 0 else x / 2 """ xp = array_namespace(x) - return xp.where(x > 0, x * 2, x / 2) + return xp.where(x > 0, x * 2, _vec_arith(xp, 4, x, 2)) diff --git a/tests/unit/test_lower.py b/tests/unit/test_lower.py index 5339926..facc6c2 100644 --- a/tests/unit/test_lower.py +++ b/tests/unit/test_lower.py @@ -52,7 +52,15 @@ def test_arith_ops_map() -> None: assert p.result == ir.BinOp( "sub", ir.Ref("x"), - ir.BinOp("div", ir.BinOp("mul", ir.Ref("y"), ir.Literal(2, "int")), ir.Literal(3, "int")), + ir.FuncCall( + "_vec_arith", + ( + ir.Ref("xp"), + ir.Literal(4, "int"), + ir.BinOp("mul", ir.Ref("y"), ir.Literal(2, "int")), + ir.Literal(3, "int"), + ), + ), ) @@ -585,3 +593,10 @@ def test_maybe_unbound_after_nested_branch() -> None: " return z" ) ) + + +@pytest.mark.parametrize("op", ["%", "//"]) +def test_unfoldable_boolean_literals_keep_runtime_arithmetic(op: str) -> None: + p = lower(f" return True {op} False") + assert isinstance(p.result, ir.FuncCall) + assert p.result.fn == "_vec_arith" diff --git a/tests/unit/test_runtime.py b/tests/unit/test_runtime.py index faba071..d8f8914 100644 --- a/tests/unit/test_runtime.py +++ b/tests/unit/test_runtime.py @@ -9,6 +9,7 @@ from __future__ import annotations import numpy as np +import pytest from array_vectorize.runtime import arith from array_vectorize.runtime import dtype @@ -180,3 +181,34 @@ def test_describe_dtype_classifies_uint64_as_float64_class() -> None: # small int literals count as the minimal dtype that fits assert dtype._describe_dtype(5) == (False, 8) assert dtype._describe_dtype(np.asarray([True])) == (False, 8) + + +@pytest.mark.parametrize("left_kind", ["int64", "uint64"]) +@pytest.mark.parametrize("right_kind", ["int64", "uint64"]) +def test_integer_division_matches_python_rounding(left_kind: str, right_kind: str) -> None: + rng = np.random.default_rng(287) + + def operands(kind: str) -> np.ndarray: + info = np.iinfo(kind) + values = rng.integers(info.min, info.max, size=128, endpoint=True, dtype=kind) + boundaries = {0, 1, info.min, info.max} + for power in range(64): + for delta in (-1, 0, 1): + value = (1 << power) + delta + if info.min <= value <= info.max: + boundaries.add(value) + if info.min <= -value <= info.max: + boundaries.add(-value) + edges = np.asarray(sorted(boundaries), dtype=kind) + rng.shuffle(edges) + return np.concatenate((values, edges)) + + left, right = operands(left_kind), operands(right_kind) + # Signed and unsigned ranges have different numbers of boundaries. + size = max(len(left), len(right)) + left, right = np.resize(left, size), np.resize(right, size) + right[right == 0] = 1 + got = arith._vec_arith(xp, int(arith.ArithOp.DIV), left, right) + expected = [int(a) / int(b) for a, b in zip(left, right, strict=True)] + np.testing.assert_array_equal(got, expected) + np.testing.assert_array_equal(np.signbit(got), np.signbit(expected))