CF Bolz-Tereick · PyPy blog · October 27, 2026
\nThe Knownbits Toy abstract domain that I wrote about earlier is great for\nall kinds of bit manipulation rewrites. However, it is not particularly good at\nrecognizing where bits move around in a sequence of instructions. It has a\nnotion of "this bit is 0 or 1 or unknown" but no notion of "this bit is the\nsame as this other bit". This prevents it from optimizing something like this:
x = ...\na = x & 0x0f\nb = x << 60\nc = b >> 60\nresult = a & cHere a and c compute the same value, so it can be optimized to result = a.\nIt is possible to do this optimization with local peephole rewrite rules too,\nbut doing that requires adding a lot of these rules, and if one is missing we\nmight miss collapsing a long sequence of operations.
In this post I want to talk about a more complicated abstract domain that is\ngreat at optimizing this kind of bit movement. It's a "bit provenance" domain.\nFor every SSA variable in a trace, it computes where each bit from the variable\ncomes from (or whether it's constant 0 or 1). In particular it can\ntrack if a bit in one variable must be the same as the bit of another variable\nat another position.
For the example above, the domain would store the following provenance\ninformation. Nothing is known about x, so its bits refer to themselves.\nUsing 64-bit integers and a logical right shift, we get this table (most\nsignificant bits first):
| operation | Bit 63 | Bit 62 | Bit 61 | Bit 60 | Bit 59 | ... | Bit 4 | Bit 3 | Bit 2 | Bit 1 | Bit 0 |
|---|---|---|---|---|---|---|---|---|---|---|---|
x = ... | x[63] | x[62] | x[61] | x[60] | x[59] | ... | x[4] | x[3] | x[2] | x[1] | x[0] |
a = x & 0x0f | 0 | 0 | 0 | 0 | 0 | ... | 0 | x[3] | x[2] | x[1] | x[0] |
b = x << 60 | x[3] | x[2] | x[1] | x[0] | 0 | ... | 0 | 0 | 0 | 0 | 0 |
c = b >> 60 | 0 | 0 | 0 | 0 | 0 | ... | 0 | x[3] | x[2] | x[1] | x[0] |
result = a & c | 0 | 0 | 0 | 0 | 0 | ... | 0 | x[3] | x[2] | x[1] | x[0] |
Then we can directly conclude that result is the same value as a.
In the rest of the post I'll describe how the domain values store this\ninformation, how the transfer functions work and how we can use the abstract\ndomain to optimize.
\nWe'll start with reusing the toy optimizer from the previous posts, in the fold for completeness’ sake.
\ntoy infrastructure
\nfrom dataclasses import dataclass\n\nINTEGER_WIDTH = 64\nINT_MASK = 2 ** INTEGER_WIDTH - 1\n\nclass Value:\n def find(self):\n raise NotImplementedError("abstract")\n\n@dataclass(eq=False)\nclass Operation(Value):\n name : str\n args : list[Value]\n\n forwarded : Value | None = None\n\n def find(self) -> Value:\n op = self\n while isinstance(op, Operation):\n next = op.forwarded\n if next is None:\n return op\n op = next\n return op\n\n def arg(self, index):\n return self.args[index].find()\n\n def make_equal_to(self, value : Value):\n self.find().forwarded = value\n\n@dataclass(eq=False)\nclass Constant(Value):\n value : object\n\n def find(self):\n return self\n\nclass Block(list):\n def __getattr__(self, opname):\n def wraparg(arg):\n if not isinstance(arg, Value):\n arg = Constant(arg)\n return arg\n def make_op(*args):\n op = Operation(opname,\n [wraparg(arg) for arg in args])\n self.append(op)\n return op\n return make_op\n\ndef bb_to_str(l : Block, varprefix : str = "var"):\n def arg_to_str(arg : Value):\n if isinstance(arg, Constant):\n return str(arg.value)\n else:\n return varnames[arg]\n\n varnames = {}\n res = []\n for index, op in enumerate(l):\n # give the operation a name used while\n # printing:\n var = f"{varprefix}{index}"\n varnames[op] = var\n arguments = ", ".join(\n arg_to_str(op.arg(i))\n for i in range(len(op.args))\n )\n strop = f"{var} = {op.name}({arguments})"\n res.append(strop)\n return "\\n".join(res)\n\ndef dce(bb: Block) -> Block:\n pure = {"getarg", "int_and", "int_or", "int_xor", "int_invert",\n "int_lshift", "int_rshift"}\n used = set()\n kept = []\n for op in reversed(bb):\n if op not in used and op.name in pure:\n continue\n kept.append(op)\n for i in range(len(op.args)):\n arg = op.arg(i)\n if isinstance(arg, Operation):\n used.add(arg)\n return Block(reversed(kept))\n\ndef test_dce():\n bb = Block()\n x = bb.getarg(0)\n live = bb.int_or(x, 1)\n dead = bb.int_xor(live, 7)\n bb.int_and(dead, 15)\n forwarded = bb.int_and(x, 0)\n forwarded.make_equal_to(Constant(0))\n bb.dummy(live)\n bb.dummy(forwarded)\n\n assert bb_to_str(dce(bb), "optvar") == """\\\noptvar0 = getarg(0)\noptvar1 = int_or(optvar0, 1)\noptvar2 = dummy(optvar1)\noptvar3 = dummy(0)"""(The dce function was adapted from the Prospero blog post.)
Then we can start implementing the domain values like this:
\n@dataclass(frozen=True)\nclass SourceBit:\n source: Operation\n index: int\n inverted: bool = False\n\n@dataclass(frozen=True)\nclass BitProvenance:\n # least significant bit first\n bits: tuple[int | SourceBit, ...]\n\n @staticmethod\n def make(bits):\n bits = tuple(bits)\n if len(bits) < INTEGER_WIDTH:\n # zero-extend, mainly for tests\n bits = bits + (0, ) * (INTEGER_WIDTH - len(bits))\n assert len(bits) == INTEGER_WIDTH\n return BitProvenance(bits)\n\n @staticmethod\n def from_constant(value):\n return BitProvenance(tuple((value >> i) & 1 for i in range(INTEGER_WIDTH)))\n\n @staticmethod\n def from_source(source: Operation):\n # nothing is known about this SSA value\n return BitProvenance(tuple(SourceBit(source, i) for i in range(INTEGER_WIDTH)))\n\n def get_constant(self) -> int | None:\n res = 0\n shift = 0\n for bit in self.bits:\n if isinstance(bit, SourceBit):\n return None\n res |= bit << shift\n shift += 1\n return resEvery SourceBit refers to a bit of an SSA Operation. For an SSA value that we\nknow nothing about, we use the from_source factory method, which returns\nsomething of the form:
x = ... # x[63] x[62] x[61] x[60] x[59] ... x[ 3] x[ 2] x[ 1] x[0]For a constant, the from_constant factory returns an abstract value where all\nthe bits are concretely known values. Constants are simply 0 or 1 and need no\nsource at all.
The inverted field of a SourceBit specifies whether the bit is a direct\ncopy of the source bit, or whether the bit is flipped.
Let's make sure we got the bit twiddling in from_constant and get_constant right by writing a Hypothesis test:
from hypothesis import given, strategies\n\nints = strategies.integers(0, 2**64-1)\n\n@given(ints)\ndef test_hypothesis_from_constant_get_constant_roundtrip(val):\n assert BitProvenance.from_constant(val).get_constant() == valThe simplest transfer function to implement for this domain is for ~, bit\ninvert. It inverts constant bits and flips the inverted field of any\nSourceBit. Here's a test:
def test_abstract_invert():\n source = Operation("getarg", [Constant(0)])\n x = BitProvenance.make([1, 0, SourceBit(source, 1, True), SourceBit(source, 0)])\n res = x.abstract_invert()\n assert res == BitProvenance.make((0, 1, SourceBit(source, 1), SourceBit(source, 0, True)) + (1, ) * 60)The implementation looks like this (the reason for the result argument will become clear a bit further down):
def invert_bit(bit):\n if isinstance(bit, int):\n return bit ^ 1\n return SourceBit(bit.source, bit.index, not bit.inverted)\n\nclass BitProvenance:\n ...\n\n def abstract_invert(self, result=None):\n return BitProvenance.make(tuple(invert_bit(bit) for bit in self.bits))The shift implementations are also quite simple as long as the shift amounts\nare constant, they just move bits around. If the shifts aren't constant, we\njust know nothing about the result.
\nIf we "know nothing", we want to use BitProvenance.from_source for the shift\noperation. Therefore the transfer functions need to get the shift Operation as\nan extra argument.
def test_lshift():\n source = Operation("getarg", [Constant(0)])\n x = BitProvenance.from_source(source) # no info\n\n # x << 7, the seven high bits get cut off\n seven = BitProvenance.from_constant(7)\n result = Operation("int_lshift", [source, Constant(7)])\n assert x.abstract_lshift(seven, result).bits == (0, ) * 7 + x.bits[:-7]\n\n # x << 65, constant 0\n big = BitProvenance.from_constant(65)\n result = Operation("int_lshift", [source, Constant(65)])\n assert x.abstract_lshift(big, result).get_constant() == 0\n\n # x << y, we know nothing\n amount = Operation("getarg", [Constant(1)])\n result = Operation("int_lshift", [source, amount])\n assert x.abstract_lshift(BitProvenance.from_source(amount), result) == BitProvenance.from_source(result)\n\n@given(ints, strategies.integers(0, 100))\ndef test_hypothesis_lshift_on_constants(val, shift):\n absval = BitProvenance.from_constant(val)\n absshift = BitProvenance.from_constant(shift)\n result = Operation("int_lshift", [Constant(val), Constant(shift)])\n assert absval.abstract_lshift(absshift, result).get_constant() == (val << shift) & INT_MASK\n\nclass BitProvenance:\n ...\n\n def abstract_lshift(self, other, result: Operation):\n const = other.get_constant()\n if const is None:\n return BitProvenance.from_source(result)\n if const >= INTEGER_WIDTH:\n # make oversized shift amounts return 0 for simplicity\n return BitProvenance.from_constant(0)\n if const == 0:\n return self\n return BitProvenance.make((0, ) * const + self.bits[:-const])Logical right shift is pretty similar, except we append the zeros at the other end:
\ndef test_rshift():\n source = Operation("getarg", [Constant(0)])\n x = BitProvenance.from_source(source) # no info\n\n # x >> 7, the seven low bits get cut off\n seven = BitProvenance.from_constant(7)\n result = Operation("int_rshift", [source, Constant(7)])\n assert x.abstract_rshift(seven, result).bits == x.bits[7:] + (0, ) * 7\n\n # x >> 65, constant 0\n big = BitProvenance.from_constant(65)\n result = Operation("int_rshift", [source, Constant(65)])\n assert x.abstract_rshift(big, result).get_constant() == 0\n\n # x >> y, we know nothing\n amount = Operation("getarg", [Constant(1)])\n result = Operation("int_rshift", [source, amount])\n assert x.abstract_rshift(BitProvenance.from_source(amount), result) == BitProvenance.from_source(result)\n\n@given(ints, strategies.integers(0, 100))\ndef test_hypothesis_rshift_on_constants(val, shift):\n absval = BitProvenance.from_constant(val)\n absshift = BitProvenance.from_constant(shift)\n result = Operation("int_rshift", [Constant(val), Constant(shift)])\n assert absval.abstract_rshift(absshift, result).get_constant() == val >> shift\n\nclass BitProvenance:\n ...\n\n def abstract_rshift(self, other, result: Operation):\n const = other.get_constant()\n if const is None:\n return BitProvenance.from_source(result)\n if const >= INTEGER_WIDTH:\n return BitProvenance.from_constant(0)\n if const == 0:\n return self\n return BitProvenance.make(self.bits[const:] + (0, ) * const)Next are the abstract binary transfer functions and, or, xor.\nWe can think about the behaviour on individual bits first. A few cases are\nsimple (e.g. a & 0 == 0) but some need reasoning about sources.
def and_bits(a, b):\n # None means the result cannot be represented as a constant or source\n # bit\n if a == 0:\n return 0\n if b == 0:\n return 0\n if a == 1:\n return b\n if b == 1:\n return a\n assert isinstance(a, SourceBit)\n assert isinstance(b, SourceBit)\n if a == b:\n # a & a == a\n return a\n if a == invert_bit(b):\n # a & not a == 0\n return 0\n return None\n\ndef test_and_bits():\n x = Operation("getarg", [Constant(0)])\n y = Operation("getarg", [Constant(1)])\n a = SourceBit(x, 3)\n assert and_bits(a, 0) == and_bits(0, a) == 0\n assert and_bits(a, 1) == and_bits(1, a) == a\n assert and_bits(a, a) == a\n assert and_bits(a, invert_bit(a)) == 0\n\n # positions disagree\n assert and_bits(a, SourceBit(x, 4)) is None\n\n # sources disagree\n assert and_bits(a, SourceBit(y, 3)) is None\n assert and_bits(a, SourceBit(y, 3, True)) is NoneWhen and_bits returns None for output bit i, we can use\nSourceBit(result, i): we don't know how that bit relates to the inputs.\nThe other output bits keep their own provenance, even if they refer to\ndifferent operations.
Now we can combine and_bits calls for all the bits of a BitProvenance.
def test_and():\n x = Operation("getarg", [Constant(0)])\n y = Operation("getarg", [Constant(1)])\n absx = BitProvenance.from_source(x)\n absy = BitProvenance.from_source(y)\n\n # x & 0 == 0\n result = Operation("int_and", [x, Constant(0)])\n assert absx.abstract_and(BitProvenance.from_constant(0), result).get_constant() == 0\n\n # x & -1 == x\n result = Operation("int_and", [x, Constant(INT_MASK)])\n assert absx.abstract_and(BitProvenance.from_constant(INT_MASK), result) == absx\n\n # x & x == x\n result = Operation("int_and", [x, x])\n assert absx.abstract_and(absx, result) == absx\n\n # x & y: we can't usefully represent provenance, so we get fresh bits\n result = Operation("int_and", [x, y])\n assert absx.abstract_and(absy, result) == BitProvenance.from_source(result)\n\n # x & ~x == 0\n result = Operation("int_and", [x, Operation("int_invert", [x])])\n assert absx.abstract_and(absx.abstract_invert(), result).get_constant() == 0\n\n@given(ints, ints)\ndef test_hypothesis_and_on_constants(a, b):\n abs_a = BitProvenance.from_constant(a)\n abs_b = BitProvenance.from_constant(b)\n result = Operation("int_and", [Constant(a), Constant(b)])\n assert abs_a.abstract_and(abs_b, result).get_constant() == a & b\n\nclass BitProvenance:\n ...\n\n def abstract_and(self, other, result: Operation):\n bits = []\n for i, (bit_a, bit_b) in enumerate(zip(self.bits, other.bits)):\n bit_res = and_bits(bit_a, bit_b)\n if bit_res is None:\n # not representable, use fresh provenance\n bit_res = SourceBit(result, i)\n bits.append(bit_res)\n return BitProvenance.make(bits)or is very similar:
def or_bits(a, b):\n if a == 1:\n return 1\n if b == 1:\n return 1\n if a == 0:\n return b\n if b == 0:\n return a\n assert isinstance(a, SourceBit)\n assert isinstance(b, SourceBit)\n if a == b:\n # a | a == a\n return a\n if a == invert_bit(b):\n # a | not a == 1\n return 1\n return None\n\ndef test_or_bits():\n x = Operation("getarg", [Constant(0)])\n y = Operation("getarg", [Constant(1)])\n a = SourceBit(x, 3)\n assert or_bits(a, 1) == or_bits(1, a) == 1\n assert or_bits(a, 0) == or_bits(0, a) == a\n assert or_bits(a, a) == a\n assert or_bits(a, invert_bit(a)) == 1\n\n # positions disagree\n assert or_bits(a, SourceBit(x, 4)) is None\n\n # sources disagree\n assert or_bits(a, SourceBit(y, 3)) is None\n assert or_bits(a, SourceBit(y, 3, True)) is None\n\ndef test_or():\n x = Operation("getarg", [Constant(0)])\n y = Operation("getarg", [Constant(1)])\n absx = BitProvenance.from_source(x)\n absy = BitProvenance.from_source(y)\n\n # x | 0 == x\n result = Operation("int_or", [x, Constant(0)])\n assert absx.abstract_or(BitProvenance.from_constant(0), result) == absx\n\n # x | -1 == -1\n result = Operation("int_or", [x, Constant(INT_MASK)])\n assert absx.abstract_or(BitProvenance.from_constant(INT_MASK), result).get_constant() == INT_MASK\n\n # x | x == x\n result = Operation("int_or", [x, x])\n assert absx.abstract_or(absx, result) == absx\n\n # x | y: we can't usefully represent provenance, so we get fresh bits\n result = Operation("int_or", [x, y])\n assert absx.abstract_or(absy, result) == BitProvenance.from_source(result)\n\n # x | ~x == -1\n result = Operation("int_or", [x, Operation("int_invert", [x])])\n assert absx.abstract_or(absx.abstract_invert(), result).get_constant() == INT_MASK\n\n@given(ints, ints)\ndef test_hypothesis_or_on_constants(a, b):\n abs_a = BitProvenance.from_constant(a)\n abs_b = BitProvenance.from_constant(b)\n result = Operation("int_or", [Constant(a), Constant(b)])\n assert abs_a.abstract_or(abs_b, result).get_constant() == a | b\n\nclass BitProvenance:\n ...\n\n def abstract_or(self, other, result: Operation):\n bits = []\n for i, (bit_a, bit_b) in enumerate(zip(self.bits, other.bits)):\n bit_res = or_bits(bit_a, bit_b)\n if bit_res is None:\n # not representable, use fresh provenance\n bit_res = SourceBit(result, i)\n bits.append(bit_res)\n return BitProvenance.make(bits)And finally here is xor:
def xor_bits(a, b):\n if a == 1:\n # x ^ 1 == ~x\n return invert_bit(b)\n if b == 1:\n return invert_bit(a)\n if a == 0:\n # x ^ 0 == x\n return b\n if b == 0:\n return a\n assert isinstance(a, SourceBit)\n assert isinstance(b, SourceBit)\n if a == b:\n # a ^ a == 0\n return 0\n if a == invert_bit(b):\n # a ^ not a == 1\n return 1\n return None\n\ndef test_xor_bits():\n x = Operation("getarg", [Constant(0)])\n y = Operation("getarg", [Constant(1)])\n a = SourceBit(x, 3)\n assert xor_bits(a, 1) == xor_bits(1, a) == invert_bit(a)\n assert xor_bits(a, 0) == xor_bits(0, a) == a\n assert xor_bits(a, a) == 0\n assert xor_bits(a, invert_bit(a)) == 1\n\n # positions disagree\n assert xor_bits(a, SourceBit(x, 4)) is None\n\n # sources disagree\n assert xor_bits(a, SourceBit(y, 3)) is None\n assert xor_bits(a, SourceBit(y, 3, True)) is None\n\ndef test_xor():\n x = Operation("getarg", [Constant(0)])\n y = Operation("getarg", [Constant(1)])\n absx = BitProvenance.from_source(x)\n absy = BitProvenance.from_source(y)\n\n # x ^ 0 == x\n result = Operation("int_xor", [x, Constant(0)])\n assert absx.abstract_xor(BitProvenance.from_constant(0), result) == absx\n\n # x ^ -1 == ~x\n result = Operation("int_xor", [x, Constant(INT_MASK)])\n assert absx.abstract_xor(BitProvenance.from_constant(INT_MASK), result) == absx.abstract_invert()\n\n # x ^ x == 0\n result = Operation("int_xor", [x, x])\n assert absx.abstract_xor(absx, result).get_constant() == 0\n\n # x ^ y: we can't usefully represent provenance, so we get fresh bits\n result = Operation("int_xor", [x, y])\n assert absx.abstract_xor(absy, result) == BitProvenance.from_source(result)\n\n # x ^ ~x == -1\n result = Operation("int_xor", [x, Operation("int_invert", [x])])\n assert absx.abstract_xor(absx.abstract_invert(), result).get_constant() == INT_MASK\n\n@given(ints, ints)\ndef test_hypothesis_xor_on_constants(a, b):\n abs_a = BitProvenance.from_constant(a)\n abs_b = BitProvenance.from_constant(b)\n result = Operation("int_xor", [Constant(a), Constant(b)])\n assert abs_a.abstract_xor(abs_b, result).get_constant() == a ^ b\n\nclass BitProvenance:\n ...\n\n def abstract_xor(self, other, result: Operation):\n bits = []\n for i, (bit_a, bit_b) in enumerate(zip(self.bits, other.bits)):\n bit_res = xor_bits(bit_a, bit_b)\n if bit_res is None:\n # not representable, use fresh provenance\n bit_res = SourceBit(result, i)\n bits.append(bit_res)\n return BitProvenance.make(bits)Now that we have the analysis, we need to find out how to use the information\nduring optimization. We'll discuss four ways of doing that, from easy to more complicated.
\nThe simplest thing to do is to exploit the fact that provenance can also\nidentify constant results.
\ndef test_constant_via_provenance():\n bb = Block()\n x = bb.getarg(0)\n a = bb.int_and(x, 0x0f)\n b = bb.int_or(x, 0xf0)\n c = bb.int_and(b, 0x0f)\n # we compute result = (x & 0x0f) ^ ((x | 0xf0) & 0x0f)\n # ((x | 0xf0) & 0x0f) can be rewritten to\n # (x & 0x0f) | (0xf0 & 0x0f)\n # which is x & 0x0f\n # this means result == (x & 0x0f) ^ (x & 0x0f) == 0\n result = bb.int_xor(a, c)\n bb.dummy(result)\n\n opt_bb = dce(simplify(bb))\n assert bb_to_str(opt_bb, "optvar") == """\\\noptvar0 = dummy(0)"""To make this test pass, we can do something very similar to what we did in the\nKnownBits post. We go along a basic block, apply the transfer functions, and if\na result is constant we can replace the operation with that constant.
\ndef unknown_transfer_functions(*abstract_args, result):\n return BitProvenance.from_source(result)\n\ndef simplify(bb: Block) -> Block:\n abstract_values = {} # dict mapping Operation to BitProvenance\n\n def provenance_of(val : Value):\n if isinstance(val, Constant):\n return BitProvenance.from_constant(val.value)\n return abstract_values[val]\n\n opt_bb = Block()\n for op in bb:\n # apply the transfer function on the abstract arguments\n name_without_prefix = op.name.removeprefix("int_")\n method_name = f"abstract_{name_without_prefix}"\n transfer_function = getattr(BitProvenance, method_name, unknown_transfer_functions)\n abstract_args = [provenance_of(arg.find()) for arg in op.args]\n abstract_res = abstract_values[op] = transfer_function(*abstract_args, result=op)\n # if the result is a constant, we optimize the operation away and make\n # it equal to the constant result\n const = abstract_res.get_constant()\n if const is not None:\n op.make_equal_to(Constant(const))\n continue\n # otherwise emit the op\n opt_bb.append(op)\n return opt_bbThe next optimization we can add is that sometimes complicated bit reshuffling\nleads to a result that was already computed earlier in the block. We can start\nwith a simple test that splits a value into low and high bits and recombines them:
\ndef test_identity_by_recombining_bits():\n # split and then recombine x\n # lo = x & 0x0f\n # hi = (x >> 4) << 4\n # result = lo | hi\n bb = Block()\n x = bb.getarg(0)\n lo = bb.int_and(x, 0x0f)\n shifted = bb.int_rshift(x, 4)\n hi = bb.int_lshift(shifted, 4)\n result = bb.int_or(lo, hi)\n bb.dummy(result)\n\n opt_bb = dce(simplify(bb))\n assert bb_to_str(opt_bb, "optvar") == """\\\noptvar0 = getarg(0)\noptvar1 = dummy(optvar0)"""To help implement this, we can add a method to check whether a provenance is\nidentical to that of an earlier result. The idea is that if the bits of a\nBitProvenance come from bits 0, 1, 2, ..., INTEGER_WIDTH - 1 of the same source,\nwithout inversion, then the operation we are looking at computes the same result\nas the source too.
def test_provenance_get_identical_source():\n source = Operation("getarg", [Constant(0)])\n x = BitProvenance.make([1, 0, SourceBit(source, 1, True), SourceBit(source, 0)])\n assert x.get_identical_source() is None # not the same value as source\n\n # trivially, a BitProvenance where we know nothing is a valid source for itself\n res = BitProvenance.from_source(source)\n assert res.get_identical_source() is source\n\nclass BitProvenance:\n ...\n def get_identical_source(self) -> Operation | None:\n bit0 = self.bits[0]\n if isinstance(bit0, int):\n return None # a constant bit cannot match an unconstrained source bit\n source = bit0.source\n # are all the bits the same as in the unconstrained source?\n if self == BitProvenance.from_source(source):\n return source\n return NoneNow we can use that to recognize identities when optimizing:
\ndef simplify(bb: Block) -> Block:\n abstract_values = {} # dict mapping Operation to BitProvenance\n\n def provenance_of(val : Value):\n if isinstance(val, Constant):\n return BitProvenance.from_constant(val.value)\n return abstract_values[val]\n\n opt_bb = Block()\n for op in bb:\n # apply the transfer function on the abstract arguments\n name_without_prefix = op.name.removeprefix("int_")\n method_name = f"abstract_{name_without_prefix}"\n transfer_function = getattr(BitProvenance, method_name, unknown_transfer_functions)\n abstract_args = [provenance_of(arg.find()) for arg in op.args]\n abstract_res = abstract_values[op] = transfer_function(*abstract_args, result=op)\n # if the result is a constant, we optimize the operation away and make\n # it equal to the constant result\n const = abstract_res.get_constant()\n if const is not None:\n op.make_equal_to(Constant(const))\n continue\n source = abstract_res.get_identical_source()\n if source is not None and source is not op:\n op.make_equal_to(source)\n continue\n # otherwise emit the op\n opt_bb.append(op)\n return opt_bbThe motivating example from the start of the blog post is next. It doesn't\nquite work with the idea from the previous subsection. Here it is in test form:
\ndef test_identity_via_provenance():\n bb = Block()\n x = bb.getarg(0)\n a = bb.int_and(x, 0x0f)\n b = bb.int_lshift(x, 60)\n c = bb.int_rshift(b, 60)\n # result is same as a, but the bits come from x\n result = bb.int_and(a, c)\n bb.dummy(result)\n\n opt_bb = dce(simplify(bb))\n assert bb_to_str(opt_bb, "optvar") == """\\\noptvar0 = getarg(0)\noptvar1 = int_and(optvar0, 15)\noptvar2 = dummy(optvar1)"""We will use a hashmap to store all the provenance patterns we have seen in the\nbasic block so far.
\ndef simplify(bb: Block) -> Block:\n abstract_values = {} # dict mapping Operation to BitProvenance\n seen_provenances = {} # dict mapping BitProvenance to Value\n\n def provenance_of(val : Value):\n if isinstance(val, Constant):\n return BitProvenance.from_constant(val.value)\n return abstract_values[val]\n\n opt_bb = Block()\n for op in bb:\n # apply the transfer function on the abstract arguments\n name_without_prefix = op.name.removeprefix("int_")\n method_name = f"abstract_{name_without_prefix}"\n transfer_function = getattr(BitProvenance, method_name, unknown_transfer_functions)\n abstract_args = [provenance_of(arg.find()) for arg in op.args]\n abstract_res = abstract_values[op] = transfer_function(*abstract_args, result=op)\n # if the result is a constant, we optimize the operation away and make\n # it equal to the constant result\n const = abstract_res.get_constant()\n if const is not None:\n op.make_equal_to(Constant(const))\n continue\n # check whether we've seen the provenance before, if yes, reuse the op\n source = seen_provenances.get(abstract_res)\n if source is not None:\n op.make_equal_to(source)\n continue\n seen_provenances[abstract_res] = op\n # otherwise emit the op\n opt_bb.append(op)\n return opt_bbIn addition to using existing values, we can sometimes emit a new shift\noperation of some earlier value, to get fewer operations than how the value was\nactually computed. Here's an example:
\ndef test_extract_rshift_from_source():\n bb = Block()\n x = bb.getarg(0)\n # lo = (x & 0xf)\n lo = bb.int_and(x, 0x0f)\n # cleared = x ^ (x & 0xf)\n cleared = bb.int_xor(x, lo)\n # result = cleared >> 4\n result = bb.int_rshift(cleared, 4)\n bb.dummy(result)\n\n # the xor clears only the bits that the shift discards\n # the result's provenance matches a direct x >> 4\n opt_bb = dce(simplify(bb))\n assert bb_to_str(opt_bb, "optvar") == """\\\noptvar0 = getarg(0)\noptvar1 = int_rshift(optvar0, 4)\noptvar2 = dummy(optvar1)"""To implement this, we add a new method BitProvenance.recognize which\ndetects whether a BitProvenance instance could be represented as a shift of\nthe source.
def test_recognize_rshift():\n x = Operation("getarg", [Constant(0)])\n # low bits come from x[4], x[5], ..., x[63]; the high four are zero\n bits = tuple(SourceBit(x, i) for i in range(4, INTEGER_WIDTH))\n shifted = BitProvenance.make(bits + (0, 0, 0, 0))\n assert shifted.recognize() == ("int_rshift", x, 4)\n\n # the same alignment is not enough if a high bit is one\n not_a_shift = BitProvenance.make(bits + (0, 0, 0, 1))\n assert not_a_shift.recognize() is None\n\n # the first bit suggests >> 4, but the next bit comes from x[6], not x[5]\n misplaced = bits[:1] + (SourceBit(x, 6),) + bits[2:]\n not_a_shift = BitProvenance.make(misplaced + (0, 0, 0, 0))\n assert not_a_shift.recognize() is None\n\n@given(strategies.integers(1, 63))\ndef test_hypothesis_recognize_shift(shift):\n x = Operation("getarg", [Constant(0)])\n absval = BitProvenance.from_source(x)\n shift_amount = BitProvenance.from_constant(shift)\n # check that lshifts are recognized\n result = Operation("int_lshift", [x, Constant(shift)])\n res = absval.abstract_lshift(shift_amount, result).recognize()\n assert res == ("int_lshift", x, shift)\n\n # check that rshifts are recognized\n result = Operation("int_rshift", [x, Constant(shift)])\n res = absval.abstract_rshift(shift_amount, result).recognize()\n assert res == ("int_rshift", x, shift)\n\nclass BitProvenance:\n ...\n def recognize(self):\n # return None or a triple (opname, source, shift amount), where\n # opname is "int_rshift" or "int_lshift"\n bits = [(i, b) for i, b in enumerate(self.bits) if isinstance(b, SourceBit)]\n if not bits:\n return None\n # try to guess the shift amount\n i, b = bits[0]\n delta = i - b.index\n source = b.source\n original = BitProvenance.from_source(source)\n # use the transfer functions to check whether it really is that shift\n if delta >= 0:\n expected = original.abstract_lshift(BitProvenance.from_constant(delta), None)\n # does the shift of the source match self?\n if self == expected:\n return ("int_lshift", source, delta)\n else:\n expected = original.abstract_rshift(BitProvenance.from_constant(-delta), None)\n # does the shift of the source match self?\n if self == expected:\n return ("int_rshift", source, -delta)\n return NoneUsing the abstract transfer functions in recognize has the nice effect that\nwe only specify what lshift and rshift do in a single place.
Now we can start using this in simplify:
def simplify(bb: Block) -> Block:\n abstract_values = {} # dict mapping Operation to BitProvenance\n seen_provenances = {} # dict mapping BitProvenance to Value\n\n def provenance_of(val : Value):\n if isinstance(val, Constant):\n return BitProvenance.from_constant(val.value)\n return abstract_values[val]\n\n opt_bb = Block()\n for op in bb:\n # apply the transfer function on the abstract arguments\n name_without_prefix = op.name.removeprefix("int_")\n method_name = f"abstract_{name_without_prefix}"\n transfer_function = getattr(BitProvenance, method_name, unknown_transfer_functions)\n abstract_args = [provenance_of(arg.find()) for arg in op.args]\n abstract_res = abstract_values[op] = transfer_function(*abstract_args, result=op)\n # if the result is a constant, we optimize the operation away and make\n # it equal to the constant result\n const = abstract_res.get_constant()\n if const is not None:\n op.make_equal_to(Constant(const))\n continue\n # check whether we've seen the provenance before, if yes, reuse the op\n source = seen_provenances.get(abstract_res)\n if source is not None:\n op.make_equal_to(source)\n continue\n\n # now try to synthesize a shift with the same provenance\n pattern = abstract_res.recognize()\n if pattern is not None:\n name, source_op, shift_amount = pattern\n if source_op is not op:\n assert shift_amount > 0 # otherwise one of the earlier rules would have caught it\n new_op = getattr(opt_bb, name)(source_op, Constant(shift_amount))\n abstract_values[new_op] = abstract_res\n seen_provenances[abstract_res] = new_op\n op.make_equal_to(new_op)\n continue\n # otherwise emit the op\n seen_provenances[abstract_res] = op\n opt_bb.append(op)\n return opt_bbThat's the whole optimizer. We propagate where each bit comes from, then look\nfor a cheaper way to produce the resulting provenance. Constant folding, reuse\nof earlier values, and shift recognition all use the same information. None of\nthem needs to know which sequence of operations produced it.
\nNow we can look at a few other rewrites that fall out of this design. Here are\nsome further passing tests:
\ndef test_xor_masks_cancel():\n bb = Block()\n x = bb.getarg(0)\n # each low-byte bit is flipped twice: 0x55 ^ 0xff ^ 0xaa == 0\n # therefore (((x ^ 0x55) ^ 0xff) ^ 0xaa) == x\n # by constant reassociation\n a = bb.int_xor(x, 0x55)\n b = bb.int_xor(a, 0xff)\n result = bb.int_xor(b, 0xaa)\n bb.dummy(result)\n\n opt_bb = dce(simplify(bb))\n assert bb_to_str(opt_bb, "optvar") == """\\\noptvar0 = getarg(0)\noptvar1 = dummy(optvar0)"""\n\ndef test_recover_low_bits_after_packing_two_sources():\n bb = Block()\n x = bb.getarg(0)\n y = bb.getarg(1)\n lo = bb.int_and(x, 0xff)\n hi = bb.int_lshift(y, 8)\n packed = bb.int_or(lo, hi)\n # masking the packed value recovers the earlier low byte of x\n result = bb.int_and(packed, 0xff)\n bb.dummy(result)\n\n opt_bb = dce(simplify(bb))\n assert bb_to_str(opt_bb, "optvar") == """\\\noptvar0 = getarg(0)\noptvar1 = int_and(optvar0, 255)\noptvar2 = dummy(optvar1)"""\n\ndef test_reuse_mask_via_de_morgan():\n bb = Block()\n x = bb.getarg(0)\n masked = bb.int_and(x, 0x0f)\n inverted = bb.int_invert(x)\n high_ones = bb.int_or(inverted, INT_MASK ^ 0x0f)\n # ~(~x | ~0x0f) == x & 0x0f, which we already computed\n result = bb.int_invert(high_ones)\n bb.dummy(masked, result)\n\n opt_bb = dce(simplify(bb))\n assert bb_to_str(opt_bb, "optvar") == """\\\noptvar0 = getarg(0)\noptvar1 = int_and(optvar0, 15)\noptvar2 = dummy(optvar1, optvar1)"""\n\ndef test_byte_swap_twice():\n def byte_swap(bb, value):\n result = 0\n for byte_index in range(8):\n byte = bb.int_and(bb.int_rshift(value, byte_index * 8), 0xff)\n moved = bb.int_lshift(byte, (7 - byte_index) * 8)\n result = bb.int_or(result, moved)\n return result\n\n bb = Block()\n x = bb.getarg(0)\n # little-endian host to network byte order, then back to host order\n network_order = byte_swap(bb, x)\n result = byte_swap(bb, network_order)\n bb.dummy(result)\n\n opt_bb = dce(simplify(bb))\n assert bb_to_str(opt_bb, "optvar") == """\\\noptvar0 = getarg(0)\noptvar1 = dummy(optvar0)"""Here's an even more complicated test. It reverses the bits of the argument,\nthen shifts left by 7, then reverses the bits of the result again. The input\nblock has 63 operations, and dce(simplify(bb)) reduces it to just three. The\nbit reversal below is adapted from Hacker's Delight, section 7-1.
def test_reverse_lshift_reverse_becomes_rshift():\n # Hacker's Delight, section 7-1: swap adjacent bits, then pairs,\n # nibbles, bytes, and larger fields. adapted from 32 to 64 bits\n def reverse_bits(bb, value):\n for shift, mask in [\n (1, 0x5555555555555555),\n (2, 0x3333333333333333),\n (4, 0x0f0f0f0f0f0f0f0f),\n (8, 0x00ff00ff00ff00ff),\n (16, 0x0000ffff0000ffff),\n (32, 0x00000000ffffffff),\n ]:\n # value = ((value & mask) << shift) | ((value >> shift) & mask)\n left = bb.int_lshift(bb.int_and(value, mask), shift)\n right = bb.int_and(bb.int_rshift(value, shift), mask)\n value = bb.int_or(left, right)\n return value\n\n bb = Block()\n x = bb.getarg(0)\n reversed_once = reverse_bits(bb, x)\n shifted = bb.int_lshift(reversed_once, 7)\n result = reverse_bits(bb, shifted)\n bb.dummy(result)\n\n opt_bb = dce(simplify(bb))\n # reversing before and after a left shift turns it into a right shift\n assert bb_to_str(opt_bb, "optvar") == """\\\noptvar0 = getarg(0)\noptvar1 = int_rshift(optvar0, 7)\noptvar2 = dummy(optvar1)"""What I like about all of these is how quite varied operation rewrites just fall\nout from the bit provenance approach, without needing complicated peephole\nrewrite patterns.
\nTo evaluate the bit provenance optimizer, I wrote a program that applies it to\nPyPy JIT traces. I didn't fully build it into the PyPy JIT because I was quite\nskeptical that it would help. Instead, I wrote an external optimization pass\nthat parses RPython JIT logs, applies the optimization and then produces\nstatistics. This uses an extended prototype that supports additional operations\nsuch as arithmetic right shifts and sign extension, and can extract shifts of\nearlier computed values as well as shifts of the original source.
\nI did this experiment for two different RPython projects, for PyPy (running the\ntranslator, that bootstraps PyPy to C code) and Pydrofoil (RISC-V simulation\nbased on Sail, booting the Linux kernel). I picked the latter because Pydrofoil\ndoes a lot of integer operations, whereas most Python code does very little.
\nFor both of these programs I wanted to know how well the bit provenance\noptimizer can optimize the traces compared to the integer optimizations of the\nPyPy JIT. I took the unoptimized traces, removed everything but integer\noperations and performed dead code elimination. To do that, the integer results\nof non-pure operations are turned into unconstrained arguments of the trace.\nInteger values used by removed operations are retained as outputs. The\nremaining integer operations are our baseline traces for this experiment.
\nThe following table shows the results. The first column shows the number of\nunoptimized integer operations. The second column shows how many operations are\nleft after running just bit provenance optimizations. The third column shows\nhow many operations are left after the RPython JIT integer optimizations run.\nThe last column shows how many operations are left after we run first the\nRPython JIT optimizer and then bit provenance optimizations. All three paths\ninclude dead code elimination after optimization.
\nIn parentheses are the percentages of how many of the unoptimized operations\nare removed by the corresponding path.
\n| Workload | Unoptimized | Provenance | RPython | RPython then Provenance |
|---|---|---|---|---|
| Pydrofoil | 1,484,985 | 822,864 (44.59%) | 769,758 (48.16%) | 656,555 (55.79%) |
| PyPy | 1,442,782 | 1,374,803 ( 4.71%) | 1,089,196 (24.51%) | 1,089,196 (24.51%) |
The 'RPython then Provenance' column shows that, for the PyPy workload,\nprovenance removes no additional operations after the RPython integer\noptimization passes in this experiment. For Pydrofoil it does remove more\noperations, but we will have to investigate whether this benefit carries over\nto the full JIT optimization pipeline.
\nHowever, the comparison between the 'Provenance' and 'RPython' column for\nPydrofoil shows that the design of the optimizer is quite good! The small\nprovenance optimizer achieves nearly as much reduction in integer operation\ncount as the complicated and well-tuned integer optimization passes of the\nRPython JIT.
\nLLVM has a very similar bit provenance tracking optimization as part of the\nHexagon backend. It also tracks 0 bits, 1 bits or references to specific bits\nfrom other SSA operations (it doesn't track negation). It is split into the\nBitTracker analysis and the HexagonBitSimplify optimization pass.
LLVM also has collectBitParts (source). Rather than running a separate\nforward analysis, it recursively walks SSA operations following provenance\n(with memoization and limited recursion depth for performance). Then there is a\ncheck whether the resulting provenance describes a byte swap or bit reversal.
In GCC this tracking happens on a byte level to be able to recognize byte swaps\nin the byte-swap recognizer. It also tracks memory addresses and offsets to\nrecognize wider loads assembled from separate byte loads, possibly with a byte\nswap. However, there's also a recent (as of writing unmerged) patch series\nto extend this to bit granularity.
\nThe interesting result here isn't that we beat an existing optimizer. It's\nthat a handful of transfer functions lets us discover rewrites across\narbitrarily long chains of bit operations, without writing a rule for each\ncombination. The evaluation suggests that existing integer optimizations\nalready cover much of this stuff. However, bit provenance gives us a compact\nway to describe a whole family of simplifications. It's quite cool to write an\noptimizer that can undo a bit reversal without ever knowing what a bit reversal\nis.
\nThe entire code can be found in this gist.
\nThanks to Joe Groff for the fantastic title. Thanks to Can Joshua Lehmann,\nAndy Wingo, Matti Picus, and Nico Rittinghaus for feedback on the draft, and to\nMax Bernstein for encouragement.
\nOriginally published at PyPy blog.
","headings":[{"level":2,"text":"Implementing the Domain","id":"implementing-the-domain"},{"level":2,"text":"Implementing Transfer Functions","id":"implementing-transfer-functions"},{"level":2,"text":"Implementing Bitwise Binary Transfer Functions","id":"implementing-bitwise-binary-transfer-functions"},{"level":2,"text":"Extracting Simpler Operations","id":"extracting-simpler-operations"},{"level":3,"text":"Constants","id":"constants"},{"level":3,"text":"Identity","id":"identity"},{"level":3,"text":"Reusing Earlier Values","id":"reusing-earlier-values"},{"level":3,"text":"Simple Shifts","id":"simple-shifts"},{"level":3,"text":"More Rewrites from the Same Rules","id":"more-rewrites-from-the-same-rules"},{"level":2,"text":"Evaluation with Pydrofoil and PyPy","id":"evaluation-with-pydrofoil-and-pypy"},{"level":2,"text":"Related work","id":"related-work"},{"level":2,"text":"Conclusion","id":"conclusion"},{"level":2,"text":"Acknowledgements","id":"acknowledgements"}]}}