Zeroize performance on u8 arrays

[blckngm]

I inspected the generated assembly code and benchmarked zeroize for [u8; 32] on x86_64 and found it quite inefficient, storing one byte at a time:

https://play.rust-lang.org/?version=stable&mode=release&edition=2021&gist=3f44f4b90e6af0eac0dcb1f649390329

On my Ryzen CPU, it takes ~7.8324 ns, or ~1cpb. Binary code size is also quite large.

Using inline assembly (just stabilized in 1.59) and SSE2, zeroing a [u8; 32] takes just 3 instructions and ~492.87 ps (~16 bytes per cycle):

let mut buf: [u8; 32];
core::arch::asm!(
    "xorps {zero}, {zero}",
    "movups {zero}, ({ptr})",
    "movups {zero}, 16({ptr})",
    zero = out(xmm_reg) _,
    ptr = in(reg) &mut buf,
    options(att_syntax, nostack, preserves_flags),
);

So it might be something worth optimizing/documenting.

If you do not want to use inline assembly, maybe you should encourage using larger types or SIMD types, e.g., [u64; 4] or [__m128; 2] instead of [u8; 32]. Using write_volatile on *mut __m128 generates equally compact and efficient code as the assembly code above.

[tarcieri]

@sopium we've done practically no optimization work, so I'm glad you're thinking about it.

Inline assembly is fine, but would need to be feature gated due to the higher MSRV (we can eventually consider bumping it, but our last bump was already painful)

For something like [T; N] though, it would be good to ensure this doesn't monomorphize to N distinct blobs of ASM.

[blckngm]

It boils down to LLVM not able to optimize write_volatile for u8 pointers. The bytes have to be written individually.

I think monomorphize for 16/32/64-byte arrays should be fine. These would also be commonly used key sizes in cryptography so the optimization can be justified.

For a u8 slice with unknown but largish size you can also use align_to_mut:

pub fn zeroize_slice(x: &mut [u8]) {
    unsafe {
        // Use __m128 on x86/x86_64.
        let (p, m, s) = x.align_to_mut::<u64>();
        for b in p {
            core::ptr::write_volatile(b, 0);
        }
        for m in m {
            core::ptr::write_volatile(m, 0);
        }
        for b in s {
            core::ptr::write_volatile(b, 0);
        }
    }
}

[tarcieri]

The use of write_volatile was something of a hack to ensure that the writes would not be elided by the compiler.

But inline assembly should also ensure that as well, and permit a highly performant implementation.

[blckngm]

The use of write_volatile was something of a hack to ensure that the writes would not be elided by the compiler.

But inline assembly should also ensure that as well, and permit a highly performant implementation.

Yes, I understand this. If it's “normal” write in a loop LLVM would be able recognize it and generate efficient instructions or call memset.

[tarcieri]

Yeah the typical way to implement a zeroization primitive efficiently is some sort of volatile memset typically realized by using some kind of compiler fence, often implemented as invoking memset in conjunction with an assembly-based dataflow optimization barrier.

While Rust has bindings to LLVM's volatile memset intrinsics, they are perma-unstable and no work has been done on a stable API AFAIK:

https://doc.rust-lang.org/std/intrinsics/fn.volatile_set_memory.html

But perhaps inline assembly could provide the next best thing, albeit on a target architecture-by-architecture basis.

[newpavlov]

We can not use write_volatile on *mut __m128 which we got from [u8; 32], since the pointer could be insufficiently aligned. AFAIK we do not have unaligned version of the volatile write. As a hack we could split [u8; 32] into aligned and unaligned parts. In theory, compiler should be able to remove the unaligned code parts if the array is sufficiently aligned.

[blckngm]

We can not use write_volatile on *mut __m128 which we got from [u8; 32], since the pointer could be insufficiently aligned. AFAIK we do not have unaligned version of the volatile write. As a hack we could split [u8; 32] into aligned and unaligned parts. In theory, compiler should be able to remove the unaligned code parts if the array is sufficiently aligned.

I was suggesting using [__m128; N] from the start, and transmuting to &[u8; N * 16] when necessary.

As for splitting to aligned and unaligned parts, that's exactly what align_to_mut does.

[blckngm]

Yeah the typical way to implement a zeroization primitive efficiently is some sort of volatile memset typically realized by using some kind of compiler fence, often implemented as invoking memset in conjunction with an assembly-based dataflow optimization barrier.

Oh this reminds me, now that we have stable inline assembly, this can be done too, and it's not architecture dependent:

pub fn zeroize(x: &mut [u8]) {
    for b in x.iter_mut() {
        *b = 0;
    }
    unsafe {
        core::arch::asm!(
            "/* {ptr} */",
            ptr = in(reg) x.as_mut_ptr(),
            options(nostack, readonly, preserves_flags),
        );
    }
}

pub fn f() {
    let mut x = [0u8; 32];
    zeroize(&mut x);
}

Playground: https://play.rust-lang.org/?version=stable&mode=release&edition=2021&gist=32fc07b8173bd2ce233eaa9bf0d8d81a

zeroize in f is optimized to two movaps but is not optimized out.

[cbeck88]

@sopium I have a question about the use of readonly in the assembly options. (This is mostly for my own understanding.)

I watched a talk by Chandler Carruth (major contributor to LLVM) a long time ago, where he described creating similar functions in order to defeat the optimizer for benchmarking purposes: https://www.youtube.com/watch?v=nXaxk27zwlk&t=2476s

static void escape(void * p) {
    asm volatile("" : : "g"(p) : "memory")
}

I realize that this is a C++ talk and not rust, but they are both going through LLVM in the end so there's a lot of overlap.

In his talk he seems to say that it's important for the optimizer to believe that
(1) This asm has side-effects that the optimizer cannot know about. In C++ that's done with "volatile" asm, and in rust I think that's done by not marking your asm pure
(2) This asm potentially clobbers all memory. So, any reads from memory after this point have to actually go back to memory and any values cached in registers are dirty.

In your block you use the readonly flag and I'm trying to understand what this does and why this is good.

1. We are telling the optimizer, this assembly is not pure, it has some "crazy unknowable side-effect" in chandler's words
2. However, we are also telling the optimizer, this assembly is readonly and doesn't write to memory
3. It also preserves flags.

So, the optimizer has to believe that

1. This assembly can read from the memory which we are supposed to have zeroes
2. This assembly has unknowable side-effects, not affecting memory, but something else (flags? some CPU state?) that it has to care about, and whatever that something is might become wrong if we don't actually write the zeroes beforehand
3. Afterwards, memory has not been clobbered, and flags have not been changed, so we do not have to re-read any cached values in registers or modify any flags.

So 3 is good because it will help the efficiency of code using this block.

But it's hard for me to understand, if we're promising the optimizer that flags are not changed, registers are not changed (only the in(reg) for ptr is read from), and memory is not changed, what is the side-effect that the compiler still has to be worried about that prevents it from dropping this block? Or the model is just that, it still has to assume that there might be something beyond this that it doesn't understand?

[newpavlov]

@cbeck88
AFAIK readonly applies only to memory of the current process, the asm block from compiler's PoV may in theory perform a syscall and it will be well in bounds of the asm! options contract.

[cbeck88]

I see, thank you

[cbeck88]

It looks to me that one barrier to actually creating a patch that does this is the need for some form of the specialization language feature.

We currently have these impl's:

/// Impl [`Zeroize`] on arrays of types that impl [`Zeroize`].
impl<Z, const N: usize> Zeroize for [Z; N]
where
    Z: Zeroize,
{
    fn zeroize(&mut self) {
        self.iter_mut().zeroize();
    }
}

utils/zeroize/src/lib.rs

Line 361 in 91c2c21

/// Impl [`Zeroize`] on arrays of types that impl [`Zeroize`].

/// Impl [`Zeroize`] on arrays of types that impl [`Zeroize`].
impl<Z, const N: usize> Zeroize for [Z; N]
where
    Z: Zeroize,
{
    fn zeroize(&mut self) {
        self.iter_mut().zeroize();
    }
}

utils/zeroize/src/lib.rs

Line 461 in 91c2c21

impl<Z> Zeroize for [Z]

If we wanted to also add impls for the special case that Z = u8, then these will be overlapping impls, and we obviously cannot remove the existing impls.

---

I had this same problem in a hashing library: https://github.com/mobilecoinfoundation/mobilecoin/blob/b07c17fd2e1946c15b3a38c30b0fb1dc4ab516e3/crypto/digestible/src/lib.rs#L218

What I decided to do there was, simply not support my trait on the primitive type u8, so that I could have my nice impl for [T] and a nice impl for [u8]. Since it's extremely rare that someone actually needs to put a single u8 in a struct, I never actually had a problem, and I was able to do this kind of optimization without waiting for specialization to land.

In your case you already shipped a stable API which supports u8 so you probably don't want to do this approach.

[tarcieri]

u8 is pretty commonly used as a carrier for masked boolean values when writing constant-time code