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):

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 Bitwise Binary Transfer Functions

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

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