=========================================== JIT Integer Optimization Peephole Rule DSL =========================================== For integer peephole optimizations in the JIT we use a domain specific language based on pattern matching that specifies how (sequences of) integer operations should be simplified. It is implemented in the directory ``jit/metainterp/optimizeopt``. This directory contains the implementation of the DSL for integer rewrites in optimizeopt. The rules are then compiled to RPython code that executes the transformations in optimizeopt. The rewrite rules are automatically proven correct with Z3 as part of the build process. This page is an introduction to how that DSL works and how to use it. Simple transformation rules ============================ The rules in the DSL specify how integer operation can be transformed into cheaper other integer operations. A rule always consists of a name, a pattern, and a target. Here's a simple rule:: add_zero: int_add(x, 0) => x The name of the rule is ``add_zero``. It matches operations in the trace of the form ``int_add(x, 0)``, where ``x`` will match anything and ``0`` will match only the constant zero. After the ``=>`` arrow is the target of the rewrite, i.e. what the operation is rewritten to, in this case ``x``. The rule language knowns which of the operations are commutative, so ``add_zero`` will also optimize ``int_add(0, x)`` to ``x``. Variables in the pattern can repeat:: sub_x_x: int_sub(x, x) => 0 This rule matches against ``int_sub`` operations where the two arguments are the same (either the same box, or the same constant). Here's a rule with a more complicated pattern:: sub_add: int_sub(int_add(x, y), y) => x This pattern matches ``int_sub`` operations, where the first argument was produced by an ``int_add`` operation. In addition, one of the arguments of the addition has to be the same as the second argument of the subtraction. The constants ``MININT``, ``MAXINT`` and ``LONG_BIT`` (which is either 32 or 64) can be used in rules, they behave like writing numbers but allow bit-width-independent formulations:: is_true_and_minint: int_is_true(int_and(x, MININT)) => int_lt(x, 0) It is also possible to have a pattern where some arguments needs to be a constant, without specifying which constant. Those patterns look like this:: sub_add_consts: int_sub(int_add(x, C1), C2) # incomplete # more goes here => int_sub(x, C) Variables in the pattern that start with a ``C`` match against constants only. However, in this current form the rule is incomplete, because the variable ``C`` that is being used in the target operation is not defined anywhere. We will see how to compute it in the next section. Computing constants and other intermediate results =================================================== Sometimes it is necessary to compute intermediate results that are used in the target operation. To do that, there can be extra assignments between the rule head and the rule target.:: sub_add_consts: int_sub(int_add(x, C1), C2) # incomplete C = C1 + C1 => int_sub(x, C) The right hand side of such an assignment is a subset of Python syntax, supporting arithmetic using ``+``, ``-``, ``*``, and certain helper functions. However, the syntax allows you to be explicit about unsignedness for some operations. E.g. ``>>u`` exists for unsigned right shifts (and I plan to add ``>u``, ``>=u``, `` x A check is followed by a boolean expression. The variables from the pattern can be used as ``IntBound`` instances in checks (and also in assignments) to find out what the abstract interpretation knows about the value of a trace variable. Here's another example:: mul_lshift: int_mul(x, int_lshift(1, y)) check y.known_ge_const(0) and y.known_le_const(LONG_BIT) => int_lshift(x, y) It expresses that ``x * (1 << y)`` can be rewritten to ``x << y`` but checks that ``y`` is known to be between ``0`` and ``LONG_BIT``. Checks and assignments can be repeated and combined with each other:: mul_pow2_const: int_mul(x, C) check C > 0 and C & (C - 1) == 0 shift = highest_bit(C) => int_lshift(x, shift) In addition to calling methods on ``IntBound`` instances, it's also possible to access their attributes, like in this rule:: and_x_c_in_range: int_and(x, C) check x.lower >= 0 and x.upper <= C & ~(C + 1) => x Rule Ordering and Liveness =================================================== The generated optimizer code will give preference to applying rules that produce a constant or a variable as a rewrite result. Only if none of those match do rules that produce new result operations get applied. For example, the rules ``sub_x_x`` and ``sub_add`` are tried before trying ``sub_add_consts``, because the former two rules optimize to a constant and a variable respectively, while the latter produces a new operation as the result. The rule ``sub_add_consts`` has a possible problem, which is that if the intermediate result of the ``int_add`` operation in the rule head is used by some other operations, then the ``sub_add_consts`` rule does not actually reduce the number of operations (and might actually make things slightly worse due to increased register pressure). However, currently it would be extremely hard to take that kind of information into account in the optimization pass of the JIT, so we optimistically apply the rules anyway. Checking rule coverage =================================================== Every rewrite rule should have at least one unit test where it triggers. To ensure this, the tests in file ``test_optimizeintbound.py`` have an assert at the end of a test run, that every rule fired at least once (this check can somethings trigger as a false positive, for example when running individual tests with ``-k``). Printing rule statistics =================================================== The JIT can print statistics about which rule fired how often in the ``jit-intbounds-stats`` :doc:`logging <../logging>` category. For example, to print it to stdout at the end of program execution, run PyPy like this:: PYPYLOG=jit-intbounds-stats:- pypy ... The output of that will look something like this:: int_add add_reassoc_consts 2514 add_zero 107008 int_sub sub_zero 31519 sub_from_zero 523 sub_x_x 3153 sub_add_consts 159 sub_add 55 sub_sub_x_c_c 1752 sub_sub_c_x_c 0 sub_xor_x_y_y 0 sub_or_x_y_y 0 int_mul mul_zero 0 mul_one 110 mul_minus_one 0 mul_pow2_const 1456 mul_lshift 0 ... Termination and Confluence ========================================= Right now there are unfortunately no checks that the rules actually rewrite operations towards "simpler" forms. There is no cost model, and also nothing that prevents you from writing a rule like this:: neg_complication: int_neg(x) # leads to infinite rewrites => int_mul(-1, x) Doing this would lead to endless rewrites if there is also another rule that turns multiplication with -1 into negation. There is also no checking for confluence__ (yet?), i.e. the property that all rewrites starting from the same input trace always lead to the same output trace, no matter in which order the rules are applied. .. __: https://en.wikipedia.org/wiki/Confluence_(abstract_rewriting) Proofs =================================================== It is very easy to write a peephole rule that is not correct in all corner cases. Therefore all the rules are proven correct with Z3 before compiled into actual JIT code, by default. When the proof fails, a (hopefully minimal) counterexample is printed. The counterexample consists of values for all the inputs that fulfil the checks, values for the intermediate expressions, and then two *different* values for the source and the target operations. E.g. if we try to add the incorrect rule:: mul_is_add: int_mul(a, b) => int_add(a, b) We get the following counterexample as output:: Could not prove correctness of rule 'mul_is_add' in line 1 counterexample given by Z3: counterexample values: a: 0 b: 1 operation int_mul(a, b) with Z3 formula a*b has counterexample result vale: 0 BUT target expression: int_add(a, b) with Z3 formula a + b has counterexample value: 1 If we add conditions, they are taken into account:: mul_is_add: int_mul(a, b) check a.known_gt_const(1) and b.known_gt_const(2) => int_add(a, b) This leads to the following counterexample:: Could not prove correctness of rule 'mul_is_add' in line 46 counterexample given by Z3: counterexample values: a: 2 b: 3 operation int_mul(a, b) with Z3 formula a*b has counterexample result vale: 6 BUT target expression: int_add(a, b) with Z3 formula a + b has counterexample value: 5 Some ``IntBound`` methods cannot be used in Z3 proofs because they have `too complex control flow`__ If that is the case, they can have Z3-equivalent formulations defined, in the ``test_z3intbound.Z3IntBound`` class (in every case this is done, it's a potential proof hole if the Z3 friendly reformulation and the real implementation differ from each other, therefore extra care is required to make very sure they are equivalent). .. __: https://pypy.org/posts/2024/08/toy-knownbits.html#cases-where-this-style-of-z3-proof-doesnt-work). If that is too hard as well, it's possible to skip the proof of individual rules by adding ``SORRY_Z3`` to its body (but we should try not to do that too often):: eq_different_knownbits: int_eq(x, y) SORRY_Z3 check x.known_ne(y) => 0 Checking for satisfiability =================================================== In addition to checking whether the rule yields a correct optimization, we also check whether the rule can ever apply. This ensures that there are *some* runtime values that would fulfil all the checks in a rule. Here's an example of a rule violating this:: never_applies: int_is_true(x) check x.known_lt_const(0) and x.known_gt_const(0) # impossible condition, always False => x Right now the error messages are not completely easy to understand, I hope to improve this later:: Rule 'never_applies' cannot ever apply in line 1 Z3 did not manage to find values for variables x such that the following condition becomes True: And(x <= x_upper, x_lower <= x, If(x_upper < 0, x_lower > 0, x_upper < 0)) Adding new rules =================================================== To add new rules (ideally motivated by `observed problems in real traces`__), the following steps should be performed: .. __: https://pypy.org/posts/2024/07/mining-jit-traces-missing-optimizations-z3.html - Add a failing test to ``test_optimizeintbound.py``. - Add the rule to ``real.rules``. - Regenerate the Python code by running ``pypy ruleopt/generate.py`` (you need the ``z3-solver`` and ``rply`` packages installed for that). - Check that ``test_optimizeintbound.py`` passes, then run the other ``optimizeopt/`` tests (in particular ``optimizeopt/test/test_z3checktests.py`` checks that the operations and expected outputs `are sensible`__). .. __: https://pypy.org/posts/2022/12/jit-bug-finding-smt-fuzzing.html DSL Implementation Notes ================================ The implementation of the DSL is done in a relatively ad-hoc manner. It is parsed using `rply`__, there's a small type checker that tries to find common problems in how the rules are written. Z3 is used via the Python API. The pattern matching RPython code is generated using an approach inspired by Luc Maranget's paper `Compiling Pattern Matching to Good Decision Trees`__, see `this blog post`__ for an approachable introduction. .. __: https://rply.readthedocs.io/ .. __: http://moscova.inria.fr/~maranget/papers/ml05e-maranget.pdf .. __: https://compiler.club/compiling-pattern-matching/