Skip to content

Evict two tree ORAM buckets at once - #1808

Open
rdeviti wants to merge 6 commits into
data61:masterfrom
rdeviti:evict-two-tree-oram-buckets
Open

rdeviti wants to merge 6 commits into
data61:masterfrom
rdeviti:evict-two-tree-oram-buckets

Conversation

@rdeviti

@rdeviti rdeviti commented Nov 12, 2026

Copy link
Copy Markdown
Contributor

I've integrated tree ORAM into a larger and more complex program, which I prefer to compile with -M rather than work out all the points where I might want to use break_point(). However, preserving memory-instruction order increases the number of communication rounds, especially in programs with many accesses at runtime-computed addresses, such as tree ORAM. So I've profiled the tree ORAM accesses and found that this happens to the two evictions in TreeORAM.evict2, as they're implemented sequentially.

However, merging them is safe with or without -M: evict picks two different buckets of the same depth, so their children differ too, and the two evictions touch disjoint buckets.

With this PR, the code evicts both buckets at once in evict2 itself, using vectors of size two, rather than leaving the merge to the compiler. Without -M, the rounds stay about the same (within 2%). With -M, accesses get faster. E.g., with 3 parties and a RecursiveORAM of 1,024 entries, this PR removed 81–84% of the slowdown. In my program, with an OptimalPackedORAM of 16,384 to 1,048,576 entries of 128 8-bit values, it removed 60–76%, depending on the ORAM size.

So I'm sending this PR in case it can help other users who're working with the -M option on. It shouldn't affect other users that don't, nor the correctness or security of tree ORAM in general. (If you prefer, I can send the bug fixes below as separate PRs, or move the code evicting both buckets at once into a conditional branch, e.g., only for -M.)

Change

evict2 gathers its two buckets into one RAM of vectors of size two, runs pop, the choice of child and add once on it, and writes both buckets back. This leaves the ORAM in the same state as evicting the buckets sequentially. Security is unchanged too: each value opened during the eviction is now a vector with one element per bucket, and each element is masked with its own randomness, as in two separate openings.

evict2 falls back to the code on master for the ORAM debug flags (debug, debug_online, crash_on_overflow and debug_ram_size), which check single buckets, and for two corner cases: ORAMs of sgf2n values, since sgf2n has no zip to interleave the two buckets, and -c sinv, whose comparisons don't support vectors. In all these cases, programs compile to the same bytecode as on master, or fail with the same error.

Note: with -M, this makes tree ORAM cheaper, so users may want to lower optimal_threshold as it also bounds the linear scan at the base of the position map. I use 256 for ORAMs of 16,384 to 1,048,576 entries and it works well for me.

Minor bug fixes

While doing this, I found two minor bugs on master that only show when the functions are called directly on vectors, as I had to do for this PR. So I'm also fixing them:

  1. floatingpoint.Trunc(a, l, m), the truncation of a by a secret number of bits m, returned wrong values for vectors on prime fields. get_bit in oram.py uses it there: for example, on rep-field, get_bit(sint([5, 10]), 2, 4) returns [1, 1] instead of [1, 0].
  2. The constant-round branch of floatingpoint.PreOR, used with -c plain and -c inv, didn't compile on vectors (AssertionError on a register size mismatch).

Both functions create their temporary registers without a size, so the registers take the global vector size, which is 1 unless the caller sets it. For instance, the vectorized sint methods set it, so sint([5, 10]) >> sint(2) works. In contrast, in oram.py, get_bit calls Trunc directly; the buckets' pop and add also call PreOR directly. With the minor fixes in this PR, both functions explicitly set the size to that of their inputs.

While testing the fix for Trunc, I found a third minor bug, which also affects master:

  1. With -O, floatingpoint.Pow2 fails to compile when its caller has set a larger vector size than its input's, e.g., when Trunc applies it to a single shift for a vector, as when evicting both buckets at once. On master, sint([5, 10, 20]).round(8, sint(2)) fails with -O on prime fields for this reason ("Mismatch of instruction and register size"). Without -O, CISC compilation compiles Pow2 separately at its input's size, so the bug doesn't show. In this PR, Pow2 also sets the size to that of its input.

None of the three fixes changes the bytecode of the 39–40 programs in Programs/Source that I could compile without arguments, with -R 64, -F 64 or -F 64 -c plain.

Testing

I'm including the test program below, in case it helps evaluate this PR. It writes, partly overwrites and reads back a tree ORAM and prints the number of wrong reads, so it checks evicting both buckets at once and the bug fixes in one go. Timer 1 gives the time and rounds of the reads. For example, Scripts/compile-run.py -E ring -M tree_oram_check 1024 128 tree 1 prints 0 of 128 reads wrong.

Programs/Source/tree_oram_check.mpc
# Writes k entries spread over a tree ORAM of n entries, overwrites every
# third, reads them back (timer 1) and prints the number of wrong reads.
# args: <n> <k> tree <values per entry>
#    or <n> <k> packed <bit lengths of at least 24, e.g. 40_40_40>
from Compiler import oram

n, k, mode, arg = int(program.args[1]), int(program.args[2]), \
    program.args[3], program.args[4]
if mode == 'tree':
    m = int(arg)
    o = oram.RecursiveORAM(n, value_type=sint, value_length=m)
else:
    bits = [int(x) for x in arg.split('_')]
    m = len(bits)
    oram.optimal_threshold = 8  # tree ORAM also for small n
    o = oram.OptimalPackedORAM(n, entry_size=tuple(bits), value_type=sint)
stride = n // k

# the m values of entry i after round r (0: write, 1: overwrite), which differ
# between entries and rounds, and from 0, the value of an empty entry
def val(i, r):
    return [i * 7 + c * 13 + r * 101 + 1 for c in range(m)]

# writes the values of round r to entry i
def write(i, r):
    o[o.index_type(i)] = tuple(sint(x) for x in val(i, r))

# round 0: writes k entries, n // k apart
@for_range(k)
def _(j):
    write(j * stride, 0)

# round 1: overwrites every third of these entries
@for_range(0, k, 3)
def _(j):
    write(j * stride, 1)

# reads the k entries back and counts those that don't match their last write
wrong = MemValue(cint(0))
start_timer(1)
@for_range(k)
def _(j):
    i = j * stride
    r = j % 3 == 0  # 1 if round 1 overwrote entry i
    got = oram.tuplify(o[o.index_type(i)])
    bad = sum(x.reveal() != cint(e) for x, e in zip(got, val(i, r)))
    wrong.iadd(bad != 0)
stop_timer(1)
print_ln('%s of %s reads wrong', wrong.read(), k)

In summary:

  • With this PR, oram_tutorial still prints 1, and all reads were correct in every case I ran, with and without -M: up to 65,536 entries, two to four parties on rings and prime fields (ring, sy-rep-ring, rep4-ring, rep-field, semi2k, semi), edaBits (-Y), -c plain, -c inv and -O.
  • With -M, the online rounds dropped by 44–56% with 64 entries, and communication didn't grow. Without -M, they changed by less than 0.1% on rings and grew by at most 2% on prime fields.
  • Since each opening during the eviction now reveals two values in a vector (one per bucket) I checked on prime fields that the two elements never share a mask, which would reveal the difference between the two buckets' values. With this PR, none of the 300 vectors that Trunc opens (64 entries, -M) shared a mask, nor did the 40,200 that PreOR opens with -c plain. As a control, I also tried a naive fix that only makes the register for the opened value big enough for the vector (instead of running all of Trunc at the vector's size): the mask then stays a single value, added to both elements, and all 300 vectors shared a mask. So the check does catch an issue if it arises.

@mkskeller

mkskeller commented Nov 13, 2026 •

Copy link
Copy Markdown
Member

Thank you for submitting this. However, the direct setting of the global vector size is out of touch with the rest of the code and error-prone due to the fact that the reset has to be called before returning. Do you think you can fix these bugs using the vectorize decorator?

@rdeviti

rdeviti commented Nov 13, 2026

Copy link
Copy Markdown
Contributor Author

Uh, right, didn't think of it. Sure, now it sets the vector sizes with library.vectorize instead. I added a small decorator that imports it at call time because library can't be imported yet when floatingpoint is loaded. The programs I tested compile to the same bytecode as before this commit, and the fixes still don't change the bytecode of the programs in Programs/Source. Let me know if you'd prefer another way.

@rdeviti

rdeviti commented Nov 13, 2026

Copy link
Copy Markdown
Contributor Author

I've also added another fix. While testing the decorator, I found that sint(5).round(8, sint([1, 2])) returned 2 instead of [2, 1] on master because round doesn't extend the single value to the size of the shifts. This extra fix works on rings and prime fields and it relies on the Trunc fix above.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants