Add aarch64 neon simd support

[Count-Count]

Currently there are only an x86_64-specific implementations using SIMD instructions.

As described in #75 I want to try to port the implementation to aarch64 using the code in /src/x86/sse42.rs as a template and the NEON intrinsics from core::arch::aarch64 (currently nightly only).

[BurntSushi]

Thanks for opening an issue!

I'm somewhat surprised to see that you want to use sse42 as a template. Did you mean sse2? I ask because sse42 uses a fairly specialized intrinsic provided by Intel. I'm not as familiar with Neon, but I'd be surprised if such a thing existed there.

Otherwise, my main concern with this is the fact that Neon intrinsics requires nightly. In general, I don't really want my crates to have code that requires nightly, even if it's an opt-in feature. In the past, I've tried this (a few times), and it almost always results in a much higher maintenance burden than I want.

I'm not quite sure how to mitigate this unfortunately. :-/

[Count-Count]

You are totally right, I meant sse2.rs, but linked the wrong file. It is still not a straight-forward translation as aarch64 neon does not have anything like _mm_movemask_epi8 afaik.

Do you think that core::arch::aarch64 will become stable eventually or will there be major changes? If no major changes are expected I would be OK with trying to implement it now and keeping it in my fork until it becomes stable.

[BurntSushi]

Not sure. My guess is that there won't be any major changes since the goal is really just to expose what the platform provides. So there isn't a ton of design work there AIUI. But I don't think the Neon functions have a champion, so it's kind of stuck.

I do somewhat suspect that the new Apple hardware will cause a lot more people to care about this though, so perhaps there is hope for stabilization sometime soon.

[digitalresistor]

I am guessing this support will be necessary before this will compile on Rust on aarch64 on macOS Big Sur on Apple Silicon M1.

[BurntSushi]

Why would it be necessary?

[digitalresistor]

I was attempting to cargo install ripgrep, and it failed to build memchr. I came looking for an issue that mentioned aarch64.

Edit:

I will fully admit that I am not sure that I have installed Rust/Cargo correctly for full arm64 support in macOS Big Sur... so there may be something else going on.

[BurntSushi]

Please file a bug. This issue should not be necessary for compiling this crate. Please include as many details as possible.

[digitalresistor]

@BurntSushi ignore my comments. I just found: https://github.com/shepmaster/rust/blob/silicon/silicon/README.md which mentions there may be issues due to some parts being x86_64 and some parts being compiled arm64.

Using RUSTC=$(rustup which rustc) $(rustup which cargo) install --verbose ripgrep

Worked, and memchr was compiled as well. I can delete my comments.

[BurntSushi]

Please don't delete comments. There's no point. And if someone else runs into the same problem, maybe a search will turn up this conversation.

Otherwise, yeah, looks like a problem in the tooling. Not memchr.

[BurntSushi]

Update: I think the best path forward here will be to use std::simd, whenever that becomes stable. std::simd is meant to be a portable SIMD abstraction. I believe memchr's use of SIMD is simple enough where it will fit easily on top of a portable API. We'll just need to benchmark it carefully to make sure we get the same codegen that we get now.

[BurntSushi]

According to this comment on irlo, aarch64 doesn't have a "move mask" operation, which is central to how the current memchr and memmem implementations work. I spent ten minutes looking at the unstable std::simd API, but couldn't figure out whether it has an equivalent portable movemask operation. (If aarch64 really doesn't have one, then I wouldn't expect the portable API to have it either, unless they can somehow emulate it efficiently. But still, it's worth checking.) So if we really can't use movemask as is from x86, then we'll need to figure something else out from aarch64.

I personally know nothing about what kinds of SIMD routines are available on aarch64. So we'll need to either look at how other folks do this (I'd look at Go's standard library or any heavily optimized libc implementation) for aarch64. Or someone who has expertise in aarch64 SIMD.

[redzic]

I'm currently working on implementing memchr for aarch64 NEON for this crate (although there is already https://github.com/ARM-software/optimized-routines which has memchr for NEON and even SVE, but maybe unrolling the loop more could be experimented with).

It is true that aarch64 doesn't have a movemask instruction. However, it does have an instruction that adds adjacent pairwise elements, and because of this, there is a really interesting trick that you can use to get a 64 bit integer that you can then count the trailing zeros on to find the matching index.

For the sake of simplicity, I'll pretend that NEON SIMD lanes are only 2 bytes wide. Basically what you can do is use cmeq (which basically functions exactly the same as SSE2's pcmpeq), then and this with 0xF00F.

Now the first byte, if it matched, will have the value 0xF0, and otherwise it will just be 0. The second byte, if it matched, will have the value 0x0F. Normally it's just 0xFF for both bytes, but it's different since we and the result with 0xF00F. Now, if we add these two elements together (which can be done efficiently in SIMD with NEON's addp instruction), we get the combined mask of these two elements. On a full 128-bit wide register, you can do this to get a 64-bit mask, count the trailing zeros, then divide the result by 4 (since the mask has 4 bits set for each element) to get the first matching index.

So far, I've written an extremely basic implementation of memchr for aarch64 based on this, which currently is effectively just a simplified version of the optimized-routines memchr (only works with lengths that are multiples of 16, fixed needle of 5 (although this is a really easy fix, but this is just an example anyway), and no unrolling, and no special handling to ensure accesses are aligned):

; x0 - pointer to src u8 array
; x1 - length of array (bytes)
function memchr, export=1
        mov             x3, x1
        mov	        w2, 0xf00f
        ; v2 = mask register
	dup	        v2.8h, w2

        movi.16b        v1, #5
L(loop):
        ldr             q0, [x0], #16
        cmeq            v0.16b, v0.16b, v1.16b
        umaxp           v3.16b, v0.16b, v0.16b
        fmov            x2, d3
        cbnz            x2, L(match_found)
        subs            x1, x1, #16
        b.ne            L(loop)
        ; no match was found
        mov             x0, -1
        ret
L(match_found):
        sub             x3, x3, x1
        and             v0.16b, v0.16b, v2.16b
        addp            v2.16b, v0.16b, v0.16b
        fmov            x0, d2
        rbit            x0, x0
        clz             x0, x0
        lsr             x0, x0, 2
        add             x0, x3, x0
        ret
endfunc

Some of the aarch64 NEON intrinsics seem to be stable as of Rust 1.59, but I'm not sure if all of the instructions needed for memchr are covered yet. Not sure if's acceptable to use assembly directly for this crate if it's the case that the intrinsics are not available on stable yet.

[BurntSushi]

@redzic That's awesome! Thank you for working on this. :-)

Some of the aarch64 NEON intrinsics seem to be stable as of Rust 1.59, but I'm not sure if all of the instructions needed for memchr are covered yet. Not sure if's acceptable to use assembly directly for this crate if it's the case that the intrinsics are not available on stable yet.

So generally speaking, it's probably a little worse than this. I generally try to keep the MSRV for memchr somewhat conservative, so even "it's stable" is usually not enough. By "conservative," I usually like to make sure it builds on Rust compilers going back 1 year, or even try to track whatever version of rustc is in Debian stable.

With that said, for high value adds like this, I can be amenable to version sniffing rustc in build.rs, and then enabling code paths automatically if the Rust version is high enough. So it seems likely that we'll want to go that route.

In principle I'm not opposed to Assembly, particularly given that we now have inline Assembly, but I do worry about my ability to maintain such a thing. I'd really really prefer to stick to intrinsics. And I think it would be reasonable to use this crate as a motivating reason to push stabilizing of necessary intrinsics if they aren't already stable.

[redzic]

Yeah that makes sense, I definitely agree that intrinsics are significantly easier to maintain than Assembly code. I only suggested it because that's probably the only good way to immediately get NEON support out there, even for very old rustc versions.

I had a quick look in the Rust stdlib, and it looks like the intrinsics for cmeq (there seem to be many, apparently for different data types) are all nightly-only at the moment.

Given that, perhaps the best route (at least for now, assuming we don't want to mess with the unstable nightly intrinsics) is to implement it via inline assembly? Although I think the best solution would be to get the intrinsics stabilized so we can just use them directly (although even then, the build.rs stuff would still be necessary for some time because of the MSRV related issues you mentioned).

Since some of the intrinsic functions in the Rust stdlib are themselves implemented with inline assembly, to get something that reads exactly like intrinsics would and that works on stable, maybe our own versions of the intrinsics could be implemented in this crate via inline assembly, and we just use those so that it still reads like intrinsic function code would.

[BurntSushi]

To be honest, I think I'd need to see it and try to understand what kind of, if any, build issues arise as a result of using it. At the very least, if it causes issues, I may need to wind up backing it out. One of the main problems here is that I have no arm hardware at all and also no experience with it.

There are other issues too, like, does arm have the equivalent of runtime CPU feature detection that we'll need to do before using certain instructions?

I think my bottom line is that I'm not 100% opposed to using inline assembly for this, but I'm pretty close to it. The only real reason I'm open to it is that I haven't tried it before, and therefore, don't have any concrete experience to fall back on.