The Shift of Theseus — Toy Bit Provenance Abstract Domain

The Knownbits Toy abstract domain that I wrote about earlier is great for all kinds of bit manipulation rewrites. However, it is not particularly good at recognizing where bits move around in a sequence of instructions. It has a notion of "this bit is 0 or 1 or unknown" but no notion of "this bit is the same as this other bit". This prevents it from optimizing something like this:

x = ...
a = x & 0x0f
b = x << 60
c = b >> 60
result = a & c

Here a and c compute the same value, so it can be optimized to result = a. It is possible to do this optimization with local peephole rewrite rules too, but doing that requires adding a lot of these rules, and if one is missing we might miss collapsing a long sequence of operations.

In this post I want to talk about a more complicated abstract domain that is great at optimizing this kind of bit movement. It's a "bit provenance" domain. For every SSA variable in a trace, it computes where each bit from the variable comes from (or whether it's constant 0 or 1). In particular it can track if a bit in one variable must be the same as the bit of another variable at another position.

For the example above, the domain would store the following provenance information. Nothing is known about x, so its bits refer to themselves. Using 64-bit integers and a logical right shift, we get this table (most significant 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 information, how the transfer functions work and how we can use the abstract domain to optimize.

Implementing the Domain¶

We'll start with reusing the toy optimizer from the previous posts, in the fold for completeness’ sake.

toy infrastructure

from dataclasses import dataclass

INTEGER_WIDTH = 64
INT_MASK = 2 ** INTEGER_WIDTH - 1

class Value:
    def find(self):
        raise NotImplementedError("abstract")


@dataclass(eq=False)
class Operation(Value):
    name : str
    args : list[Value]

    forwarded : Value | None = None

    def find(self) -> Value:
        op = self
        while isinstance(op, Operation):
            next = op.forwarded
            if next is None:
                return op
            op = next
        return op

    def arg(self, index):
        return self.args[index].find()

    def make_equal_to(self, value : Value):
        self.find().forwarded = value


@dataclass(eq=False)
class Constant(Value):
    value : object

    def find(self):
        return self


class Block(list):
    def __getattr__(self, opname):
        def wraparg(arg):
            if not isinstance(arg, Value):
                arg = Constant(arg)
            return arg
        def make_op(*args):
            op = Operation(opname,
                [wraparg(arg) for arg in args])
            self.append(op)
            return op
        return make_op


def bb_to_str(l : Block, varprefix : str = "var"):
    def arg_to_str(arg : Value):
        if isinstance(arg, Constant):
            return str(arg.value)
        else:
            return varnames[arg]

    varnames = {}
    res = []
    for index, op in enumerate(l):
        # give the operation a name used while
        # printing:
        var =  f"{varprefix}{index}"
        varnames[op] = var
        arguments = ", ".join(
            arg_to_str(op.arg(i))
                for i in range(len(op.args))
        )
        strop = f"{var} = {op.name}({arguments})"
        res.append(strop)
    return "\n".join(res)

def dce(bb: Block) -> Block:
    pure = {"getarg", "int_and", "int_or", "int_xor", "int_invert",
            "int_lshift", "int_rshift"}
    used = set()
    kept = []
    for op in reversed(bb):
        if op not in used and op.name in pure:
            continue
        kept.append(op)
        for i in range(len(op.args)):
            arg = op.arg(i)
            if isinstance(arg, Operation):
                used.add(arg)
    return Block(reversed(kept))


def test_dce():
    bb = Block()
    x = bb.getarg(0)
    live = bb.int_or(x, 1)
    dead = bb.int_xor(live, 7)
    bb.int_and(dead, 15)
    forwarded = bb.int_and(x, 0)
    forwarded.make_equal_to(Constant(0))
    bb.dummy(live)
    bb.dummy(forwarded)

    assert bb_to_str(dce(bb), "optvar") == """\
optvar0 = getarg(0)
optvar1 = int_or(optvar0, 1)
optvar2 = dummy(optvar1)
optvar3 = dummy(0)"""

(The dce function was adapted from the Prospero blog post.)

Then we can start implementing the domain values like this:

@dataclass(frozen=True)
class SourceBit:
    source: Operation
    index: int
    inverted: bool = False


@dataclass(frozen=True)
class BitProvenance:
    # least significant bit first
    bits: tuple[int | SourceBit, ...]

    @staticmethod
    def make(bits):
        bits = tuple(bits)
        if len(bits) < INTEGER_WIDTH:
            # zero-extend, mainly for tests
            bits = bits + (0, ) * (INTEGER_WIDTH - len(bits))
        assert len(bits) == INTEGER_WIDTH
        return BitProvenance(bits)

    @staticmethod
    def from_constant(value):
        return BitProvenance(tuple((value >> i) & 1 for i in range(INTEGER_WIDTH)))

    @staticmethod
    def from_source(source: Operation):
        # nothing is known about this SSA value
        return BitProvenance(tuple(SourceBit(source, i) for i in range(INTEGER_WIDTH)))

    def get_constant(self) -> int | None:
        res = 0
        shift = 0
        for bit in self.bits:
            if isinstance(bit, SourceBit):
                return None
            res |= bit << shift
            shift += 1
        return res

Every SourceBit refers to a bit of an SSA Operation. For an SSA value that we know nothing about, we use the from_source factory method, which returns something 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 the bits are concretely known values. Constants are simply 0 or 1 and need no source at all.

The inverted field of a SourceBit specifies whether the bit is a direct copy 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

ints = strategies.integers(0, 2**64-1)

@given(ints)
def test_hypothesis_from_constant_get_constant_roundtrip(val):
    assert BitProvenance.from_constant(val).get_constant() == val

Implementing Transfer Functions¶

The simplest transfer function to implement for this domain is for ~, bit invert. It inverts constant bits and flips the inverted field of any SourceBit. Here's a test:

def test_abstract_invert():
    source = Operation("getarg", [Constant(0)])
    x = BitProvenance.make([1, 0, SourceBit(source, 1, True), SourceBit(source, 0)])
    res = x.abstract_invert()
    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):
    if isinstance(bit, int):
        return bit ^ 1
    return SourceBit(bit.source, bit.index, not bit.inverted)


class BitProvenance:
    ...

    def abstract_invert(self, result=None):
        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 are constant, they just move bits around. If the shifts aren't constant, we just know nothing about the result.

If we "know nothing", we want to use BitProvenance.from_source for the shift operation. Therefore the transfer functions need to get the shift Operation as an extra argument.

def test_lshift():
    source = Operation("getarg", [Constant(0)])
    x = BitProvenance.from_source(source) # no info

    # x << 7, the seven high bits get cut off
    seven = BitProvenance.from_constant(7)
    result = Operation("int_lshift", [source, Constant(7)])
    assert x.abstract_lshift(seven, result).bits == (0, ) * 7 + x.bits[:-7]

    # x << 65, constant 0
    big = BitProvenance.from_constant(65)
    result = Operation("int_lshift", [source, Constant(65)])
    assert x.abstract_lshift(big, result).get_constant() == 0

    # x << y, we know nothing
    amount = Operation("getarg", [Constant(1)])
    result = Operation("int_lshift", [source, amount])
    assert x.abstract_lshift(BitProvenance.from_source(amount), result) == BitProvenance.from_source(result)


@given(ints, strategies.integers(0, 100))
def test_hypothesis_lshift_on_constants(val, shift):
    absval = BitProvenance.from_constant(val)
    absshift = BitProvenance.from_constant(shift)
    result = Operation("int_lshift", [Constant(val), Constant(shift)])
    assert absval.abstract_lshift(absshift, result).get_constant() == (val << shift) & INT_MASK


class BitProvenance:
    ...

    def abstract_lshift(self, other, result: Operation):
        const = other.get_constant()
        if const is None:
            return BitProvenance.from_source(result)
        if const >= INTEGER_WIDTH:
            # make oversized shift amounts return 0 for simplicity
            return BitProvenance.from_constant(0)
        if const == 0:
            return self
        return BitProvenance.make((0, ) * const + self.bits[:-const])

Logical right shift is pretty similar, except we append the zeros at the other end:

def test_rshift():
    source = Operation("getarg", [Constant(0)])
    x = BitProvenance.from_source(source) # no info

    # x >> 7, the seven low bits get cut off
    seven = BitProvenance.from_constant(7)
    result = Operation("int_rshift", [source, Constant(7)])
    assert x.abstract_rshift(seven, result).bits == x.bits[7:] + (0, ) * 7

    # x >> 65, constant 0
    big = BitProvenance.from_constant(65)
    result = Operation("int_rshift", [source, Constant(65)])
    assert x.abstract_rshift(big, result).get_constant() == 0

    # x >> y, we know nothing
    amount = Operation("getarg", [Constant(1)])
    result = Operation("int_rshift", [source, amount])
    assert x.abstract_rshift(BitProvenance.from_source(amount), result) == BitProvenance.from_source(result)


@given(ints, strategies.integers(0, 100))
def test_hypothesis_rshift_on_constants(val, shift):
    absval = BitProvenance.from_constant(val)
    absshift = BitProvenance.from_constant(shift)
    result = Operation("int_rshift", [Constant(val), Constant(shift)])
    assert absval.abstract_rshift(absshift, result).get_constant() == val >> shift


class BitProvenance:
    ...

    def abstract_rshift(self, other, result: Operation):
        const = other.get_constant()
        if const is None:
            return BitProvenance.from_source(result)
        if const >= INTEGER_WIDTH:
            return BitProvenance.from_constant(0)
        if const == 0:
            return self
        return BitProvenance.make(self.bits[const:] + (0, ) * const)

Implementing Bitwise Binary Transfer Functions¶

Next are the abstract binary transfer functions and, or, xor. We can think about the behaviour on individual bits first. A few cases are simple (e.g. a & 0 == 0) but some need reasoning about sources.

def and_bits(a, b):
    # None means the result cannot be represented as a constant or source
    # bit
    if a == 0:
        return 0
    if b == 0:
        return 0
    if a == 1:
        return b
    if b == 1:
        return a
    assert isinstance(a, SourceBit)
    assert isinstance(b, SourceBit)
    if a == b:
        # a & a == a
        return a
    if a == invert_bit(b):
        # a & not a == 0
        return 0
    return None


def test_and_bits():
    x = Operation("getarg", [Constant(0)])
    y = Operation("getarg", [Constant(1)])
    a = SourceBit(x, 3)
    assert and_bits(a, 0) == and_bits(0, a) == 0
    assert and_bits(a, 1) == and_bits(1, a) == a
    assert and_bits(a, a) == a
    assert and_bits(a, invert_bit(a)) == 0

    # positions disagree
    assert and_bits(a, SourceBit(x, 4)) is None

    # sources disagree
    assert and_bits(a, SourceBit(y, 3)) is None
    assert and_bits(a, SourceBit(y, 3, True)) is None

When and_bits returns None for output bit i, we can use SourceBit(result, i): we don't know how that bit relates to the inputs. The other output bits keep their own provenance, even if they refer to different operations.

Now we can combine and_bits calls for all the bits of a BitProvenance.

def test_and():
    x = Operation("getarg", [Constant(0)])
    y = Operation("getarg", [Constant(1)])
    absx = BitProvenance.from_source(x)
    absy = BitProvenance.from_source(y)

    # x & 0 == 0
    result = Operation("int_and", [x, Constant(0)])
    assert absx.abstract_and(BitProvenance.from_constant(0), result).get_constant() == 0

    # x & -1 == x
    result = Operation("int_and", [x, Constant(INT_MASK)])
    assert absx.abstract_and(BitProvenance.from_constant(INT_MASK), result) == absx

    # x & x == x
    result = Operation("int_and", [x, x])
    assert absx.abstract_and(absx, result) == absx

    # x & y: we can't usefully represent provenance, so we get fresh bits
    result = Operation("int_and", [x, y])
    assert absx.abstract_and(absy, result) == BitProvenance.from_source(result)

    # x & ~x == 0
    result = Operation("int_and", [x, Operation("int_invert", [x])])
    assert absx.abstract_and(absx.abstract_invert(), result).get_constant() == 0


@given(ints, ints)
def test_hypothesis_and_on_constants(a, b):
    abs_a = BitProvenance.from_constant(a)
    abs_b = BitProvenance.from_constant(b)
    result = Operation("int_and", [Constant(a), Constant(b)])
    assert abs_a.abstract_and(abs_b, result).get_constant() == a & b


class BitProvenance:
    ...

    def abstract_and(self, other, result: Operation):
        bits = []
        for i, (bit_a, bit_b) in enumerate(zip(self.bits, other.bits)):
            bit_res = and_bits(bit_a, bit_b)
            if bit_res is None:
                # not representable, use fresh provenance
                bit_res = SourceBit(result, i)
            bits.append(bit_res)
        return BitProvenance.make(bits)

or is very similar:

def or_bits(a, b):
    if a == 1:
        return 1
    if b == 1:
        return 1
    if a == 0:
        return b
    if b == 0:
        return a
    assert isinstance(a, SourceBit)
    assert isinstance(b, SourceBit)
    if a == b:
        # a | a == a
        return a
    if a == invert_bit(b):
        # a | not a == 1
        return 1
    return None

def test_or_bits():
    x = Operation("getarg", [Constant(0)])
    y = Operation("getarg", [Constant(1)])
    a = SourceBit(x, 3)
    assert or_bits(a, 1) == or_bits(1, a) == 1
    assert or_bits(a, 0) == or_bits(0, a) == a
    assert or_bits(a, a) == a
    assert or_bits(a, invert_bit(a)) == 1

    # positions disagree
    assert or_bits(a, SourceBit(x, 4)) is None

    # sources disagree
    assert or_bits(a, SourceBit(y, 3)) is None
    assert or_bits(a, SourceBit(y, 3, True)) is None


def test_or():
    x = Operation("getarg", [Constant(0)])
    y = Operation("getarg", [Constant(1)])
    absx = BitProvenance.from_source(x)
    absy = BitProvenance.from_source(y)

    # x | 0 == x
    result = Operation("int_or", [x, Constant(0)])
    assert absx.abstract_or(BitProvenance.from_constant(0), result) == absx

    # x | -1 == -1
    result = Operation("int_or", [x, Constant(INT_MASK)])
    assert absx.abstract_or(BitProvenance.from_constant(INT_MASK), result).get_constant() == INT_MASK

    # x | x == x
    result = Operation("int_or", [x, x])
    assert absx.abstract_or(absx, result) == absx

    # x | y: we can't usefully represent provenance, so we get fresh bits
    result = Operation("int_or", [x, y])
    assert absx.abstract_or(absy, result) == BitProvenance.from_source(result)

    # x | ~x == -1
    result = Operation("int_or", [x, Operation("int_invert", [x])])
    assert absx.abstract_or(absx.abstract_invert(), result).get_constant() == INT_MASK


@given(ints, ints)
def test_hypothesis_or_on_constants(a, b):
    abs_a = BitProvenance.from_constant(a)
    abs_b = BitProvenance.from_constant(b)
    result = Operation("int_or", [Constant(a), Constant(b)])
    assert abs_a.abstract_or(abs_b, result).get_constant() == a | b


class BitProvenance:
    ...

    def abstract_or(self, other, result: Operation):
        bits = []
        for i, (bit_a, bit_b) in enumerate(zip(self.bits, other.bits)):
            bit_res = or_bits(bit_a, bit_b)
            if bit_res is None:
                # not representable, use fresh provenance
                bit_res = SourceBit(result, i)
            bits.append(bit_res)
        return BitProvenance.make(bits)

And finally here is xor:

def xor_bits(a, b):
    if a == 1:
        # x ^ 1 == ~x
        return invert_bit(b)
    if b == 1:
        return invert_bit(a)
    if a == 0:
        # x ^ 0 == x
        return b
    if b == 0:
        return a
    assert isinstance(a, SourceBit)
    assert isinstance(b, SourceBit)
    if a == b:
        # a ^ a == 0
        return 0
    if a == invert_bit(b):
        # a ^ not a == 1
        return 1
    return None

def test_xor_bits():
    x = Operation("getarg", [Constant(0)])
    y = Operation("getarg", [Constant(1)])
    a = SourceBit(x, 3)
    assert xor_bits(a, 1) == xor_bits(1, a) == invert_bit(a)
    assert xor_bits(a, 0) == xor_bits(0, a) == a
    assert xor_bits(a, a) == 0
    assert xor_bits(a, invert_bit(a)) == 1

    # positions disagree
    assert xor_bits(a, SourceBit(x, 4)) is None

    # sources disagree
    assert xor_bits(a, SourceBit(y, 3)) is None
    assert xor_bits(a, SourceBit(y, 3, True)) is None


def test_xor():
    x = Operation("getarg", [Constant(0)])
    y = Operation("getarg", [Constant(1)])
    absx = BitProvenance.from_source(x)
    absy = BitProvenance.from_source(y)

    # x ^ 0 == x
    result = Operation("int_xor", [x, Constant(0)])
    assert absx.abstract_xor(BitProvenance.from_constant(0), result) == absx

    # x ^ -1 == ~x
    result = Operation("int_xor", [x, Constant(INT_MASK)])
    assert absx.abstract_xor(BitProvenance.from_constant(INT_MASK), result) == absx.abstract_invert()

    # x ^ x == 0
    result = Operation("int_xor", [x, x])
    assert absx.abstract_xor(absx, result).get_constant() == 0

    # x ^ y: we can't usefully represent provenance, so we get fresh bits
    result = Operation("int_xor", [x, y])
    assert absx.abstract_xor(absy, result) == BitProvenance.from_source(result)

    # x ^ ~x == -1
    result = Operation("int_xor", [x, Operation("int_invert", [x])])
    assert absx.abstract_xor(absx.abstract_invert(), result).get_constant() == INT_MASK


@given(ints, ints)
def test_hypothesis_xor_on_constants(a, b):
    abs_a = BitProvenance.from_constant(a)
    abs_b = BitProvenance.from_constant(b)
    result = Operation("int_xor", [Constant(a), Constant(b)])
    assert abs_a.abstract_xor(abs_b, result).get_constant() == a ^ b


class BitProvenance:
    ...

    def abstract_xor(self, other, result: Operation):
        bits = []
        for i, (bit_a, bit_b) in enumerate(zip(self.bits, other.bits)):
            bit_res = xor_bits(bit_a, bit_b)
            if bit_res is None:
                # not representable, use fresh provenance
                bit_res = SourceBit(result, i)
            bits.append(bit_res)
        return BitProvenance.make(bits)

Evaluation with Pydrofoil and PyPy¶

To evaluate the bit provenance optimizer, I wrote a program that applies it to PyPy JIT traces. I didn't fully build it into the PyPy JIT because I was quite skeptical that it would help. Instead, I wrote an external optimization pass that parses RPython JIT logs, applies the optimization and then produces statistics. This uses an extended prototype that supports additional operations such as arithmetic right shifts and sign extension, and can extract shifts of earlier computed values as well as shifts of the original source.

I did this experiment for two different RPython projects, for PyPy (running the translator, that bootstraps PyPy to C code) and Pydrofoil (RISC-V simulation based on Sail, booting the Linux kernel). I picked the latter because Pydrofoil does a lot of integer operations, whereas most Python code does very little.

For both of these programs I wanted to know how well the bit provenance optimizer can optimize the traces compared to the integer optimizations of the PyPy JIT. I took the unoptimized traces, removed everything but integer operations and performed dead code elimination. To do that, the integer results of non-pure operations are turned into unconstrained arguments of the trace. Integer values used by removed operations are retained as outputs. The remaining integer operations are our baseline traces for this experiment.

The following table shows the results. The first column shows the number of unoptimized integer operations. The second column shows how many operations are left after running just bit provenance optimizations. The third column shows how many operations are left after the RPython JIT integer optimizations run. The last column shows how many operations are left after we run first the RPython JIT optimizer and then bit provenance optimizations. All three paths include dead code elimination after optimization.

In parentheses are the percentages of how many of the unoptimized operations are removed by the corresponding path.

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, provenance removes no additional operations after the RPython integer optimization passes in this experiment. For Pydrofoil it does remove more operations, but we will have to investigate whether this benefit carries over to the full JIT optimization pipeline.

However, the comparison between the 'Provenance' and 'RPython' column for Pydrofoil shows that the design of the optimizer is quite good! The small provenance optimizer achieves nearly as much reduction in integer operation count as the complicated and well-tuned integer optimization passes of the RPython JIT.

Conclusion¶

The interesting result here isn't that we beat an existing optimizer. It's that a handful of transfer functions lets us discover rewrites across arbitrarily long chains of bit operations, without writing a rule for each combination. The evaluation suggests that existing integer optimizations already cover much of this stuff. However, bit provenance gives us a compact way to describe a whole family of simplifications. It's quite cool to write an optimizer that can undo a bit reversal without ever knowing what a bit reversal is.

The entire code can be found in this gist.

Acknowledgements¶

Thanks to Joe Groff for the fantastic title. Thanks to Can Joshua Lehmann, Andy Wingo, Matti Picus, and Nico Rittinghaus for feedback on the draft, and to Max Bernstein for encouragement.