Re: [PATCH v5 14/15] rust: sync: reduce stack usage of `UniqueArc::try_new_uninit`

y86-dev
2023-04-03T21:07:18Z

On 03.04.23 19:56, Alice Ryhl wrote:
> On 4/3/23 18:05, Benno Lossin wrote:
>> /// Tries to allocate a new [`UniqueArc`] instance whose contents are not initialised yet.
>> pub fn try_new_uninit() -> Result<UniqueArc<MaybeUninit<T>>, AllocError> {
>> - Ok(UniqueArc::<MaybeUninit<T>> {
>> + // INVARIANT: The refcount is initialised to a non-zero value.
>> + let inner = Box::try_init::<AllocError>(try_init!(ArcInner {
>> + // SAFETY: There are no safety requirements for this FFI call.
>> + refcount: Opaque::new(unsafe { bindings::REFCOUNT_INIT(1) }),
>> + data <- init::uninit::<T, AllocError>(),
>> + }? AllocError))?;
>> + Ok(UniqueArc {
>> // INVARIANT: The newly-created object has a ref-count of 1.
>> - inner: Arc::try_new(MaybeUninit::uninit())?,
>> + // SAFETY: The pointer from the `Box` is valid.
>> + inner: unsafe { Arc::from_inner(Box::leak(inner).into()) },
>> })
>> }
>> }
>
> I'm curious - do you know whether this compiles to the same machine code
> as this?
>
> pub fn try_new_uninit() -> Result<UniqueArc<MaybeUninit<T>>, AllocError> {
> let inner: Box<MaybeUninit<ArcInner<T>>> = Box::try_new_uninit()?;
> let ptr = Box::into_raw(inner) as *mut ArcInner<T>;
> addr_of_mut!((*ptr).refcount).write(bindings::REFCOUNT_INIT(1));
> Ok(UniqueArc {
> inner: Arc {
> ptr: unsafe { NonNull::new_unchecked(ptr) },
> _p: PhantomData,
> }
> })
> }

Yes it compiles to exactly the same assembly (byte for byte), because I was
not sure if I was compiling the right thing, I compiled this function:

unsafe { core::arch::asm!("ud2") };
let r1 = UniqueArc::try_new_uninit();
unsafe { core::arch::asm!("ud2") };
let r2 = UniqueArc::try_new_uninit2();
unsafe { core::arch::asm!("ud2") };

The `ud2` instructions are for better visibility, as I have not read a lot
of assembly. The above disassembles to this:

ffffffff8143bb80 <_RNvXCsdVu6umiBwhr_12rust_minimalNtB2_11RustMinimalNtCsfATHBUcknU9_6kernel6Module4init>:
ffffffff8143bb80: 41 57 push r15
ffffffff8143bb82: 41 56 push r14
ffffffff8143bb84: 53 push rbx
ffffffff8143bb85: 49 89 ff mov r15,rdi
ffffffff8143bb88: 0f 0b ud2
ffffffff8143bb8a: bf 04 00 10 00 mov edi,0x100004
ffffffff8143bb8f: be 04 00 00 00 mov esi,0x4
ffffffff8143bb94: e8 a7 23 f5 ff call ffffffff8138df40 <__rust_alloc>
ffffffff8143bb99: 48 85 c0 test rax,rax
ffffffff8143bb9c: 74 12 je ffffffff8143bbb0 <_RNvXCsdVu6umiBwhr_12rust_minimalNtB2_11RustMinimalNtCsfATHBUcknU9_6kernel6Module4init+0x30>
ffffffff8143bb9e: 49 89 c6 mov r14,rax
ffffffff8143bba1: bf 01 00 00 00 mov edi,0x1
ffffffff8143bba6: e8 e5 fc f4 ff call ffffffff8138b890 <rust_helper_REFCOUNT_INIT>
ffffffff8143bbab: 41 89 06 mov DWORD PTR [r14],eax
ffffffff8143bbae: eb 03 jmp ffffffff8143bbb3 <_RNvXCsdVu6umiBwhr_12rust_minimalNtB2_11RustMinimalNtCsfATHBUcknU9_6kernel6Module4init+0x33>
ffffffff8143bbb0: 45 31 f6 xor r14d,r14d
ffffffff8143bbb3: 0f 0b ud2
ffffffff8143bbb5: bf 04 00 10 00 mov edi,0x100004
ffffffff8143bbba: be 04 00 00 00 mov esi,0x4
ffffffff8143bbbf: e8 7c 23 f5 ff call ffffffff8138df40 <__rust_alloc>
ffffffff8143bbc4: 48 85 c0 test rax,rax
ffffffff8143bbc7: 74 11 je ffffffff8143bbda <_RNvXCsdVu6umiBwhr_12rust_minimalNtB2_11RustMinimalNtCsfATHBUcknU9_6kernel6Module4init+0x5a>
ffffffff8143bbc9: 48 89 c3 mov rbx,rax
ffffffff8143bbcc: bf 01 00 00 00 mov edi,0x1
ffffffff8143bbd1: e8 ba fc f4 ff call ffffffff8138b890 <rust_helper_REFCOUNT_INIT>
ffffffff8143bbd6: 89 03 mov DWORD PTR [rbx],eax
ffffffff8143bbd8: eb 02 jmp ffffffff8143bbdc <_RNvXCsdVu6umiBwhr_12rust_minimalNtB2_11RustMinimalNtCsfATHBUcknU9_6kernel6Module4init+0x5c>
ffffffff8143bbda: 31 db xor ebx,ebx
ffffffff8143bbdc: 0f 0b ud2
ffffffff8143bbde: 4d 89 77 08 mov QWORD PTR [r15+0x8],r14
ffffffff8143bbe2: 49 89 5f 10 mov QWORD PTR [r15+0x10],rbx
ffffffff8143bbe6: 41 c7 07 00 00 00 00 mov DWORD PTR [r15],0x0
ffffffff8143bbed: 4c 89 f8 mov rax,r15
ffffffff8143bbf0: 5b pop rbx
ffffffff8143bbf1: 41 5e pop r14
ffffffff8143bbf3: 41 5f pop r15
ffffffff8143bbf5: c3 ret

I have not done extensive enough tests to be sure about this in general,
but most of the examples of the pin-init API that I looked at were
optimized to the same assembly as manual initialization.

The only examples were this was not the case was when I had triply nested
structs with `Box`es that all were initialized via `PinInit`. That was the
point were the initialization closure was not inlined any more.
I also verified that in this particular case the closure was again inlined
after adding `#[inline]` to it (which requires `stmt_expr_attributes`).

At some point I might do a more thorough analysis.

--
Cheers,
Benno