I’d like to propose introducing two new intrinsics to enable vectorization of loops with early
exits that have potentially faulting loads.
The new intrinsics are:
@llvm.speculative.load (name subject to change) - perform a load that
may access memory beyond the allocated object. It must be used in
combination with @llvm.can.load.speculatively to ensure the load is
guaranteed to not trap.
@llvm.can.load.speculatively - Returns true if it’s safe to speculatively
load a given number of bytes from a pointer. The semantics are
target-dependent. On some targets, this may check that the access
does not cross page boundaries, or stricter checks for example on
AArch64 with MTE, which limits the access size to 16 bytes.
@llvm.speculative.load is lowered to a regular load in SelectionDAG
without MODereferenceable. I am not sure if we need to be more careful
than this, i.e. if we could still reason about SelectionDAG loads to
infer dereferencability for the pointer.
@llvm.can.load.speculatively is lowered to regular IR in PreISel
lowering, using a target-lowering hook. By default, it conservatively
expands to false.
These intrinsics should allow the loop vectorizer to vectorize early-exit
loops with potentially non-dereferenceable loads.
This has previously been discussed in https://github.com/llvm/llvm-project/pull/120603 and is similar to @nikic’s sketch [RFC] Intrinsic for non-trapping, partially out-of-bounds read - HackMD , with the major difference
being that there is no %defined_size argument and instead the load
returns the stored values for the bytes within bounds and undef
otherwise. I don’t think we can easily compute the defined size because
it may depend on the loaded values (i.e. at what lane the early exit has
been taken).
Thanks for the proposal. I agree that we need some kind of solution in this space.
I don’t think that there should be an explicitly requirement to check this intrinsic. It should be fine to use it directly based on target knowledge. llvm.can.load.speculatively is just a way to guarantee that.
I wasn’t a fan of the %defined_size parameter either, which is why I never submitted this RFC. The reason why the parameter was there (as outlined in this section), is that I was struggling to define the semantics of the intrinsic in the presence of noalias without it.
You are defining it in terms of reads past the end of the allocated object, in which case everything is simple. But this is not the only constraint on which part of the object is accessible.
With noalias, which parts of the @llvm.speculative.load result are undef can depend on which memory is being written through noalias pointers later in the function.
Though probably we can get away with just ignoring this… at least I can’t come up with any specific issues this would cause, as long as we treat the speculative load as an upperBound access for AA purposes.
I believe that on the Rust side, there was some interest with having speculative loads with an atomic memory ordering as well, though I don’t recall what they were needed for.
How does this interact with bare-metal environments with physical memory protection that can in general have much finer granularity than a page? For example, on RISC-V, its Physical Memory Protection (“PMP”) feature can configure arbitrary regions with a 4-byte granularity (albeit less efficiently than for power of two size and alignment), and I don’t see a way for LLVM to know whether that’s been done or not.
I would expect @llvm.can.load.speculatively to expand to false on such platforms. I would also probably expect a TTI hook to determine whether or not they are generated in the first place given an expansion to a constant in PreISelIntrinsicLowering is unlikely to be properly cleaned up.
How is vectorizer expected to use that? Call @llvm.can.load.speculativelyand jump to scalar remainder if it ever returns false on any vector iteration?
Transforms wanting to create the speculative loads would check the intrinsic cost via TTI, and that should return invalid cost if it gets lowered to known false (will share patch soon)
As a first step yes (will share patch soon). In the future, we may be able to extend it so llvm.can.load.speculatively could possibly return the size that is safe to load speculatively, which could then either be used within the vector loop or peel off iterations at the beginning so that the loads in the vector loop become safe.
we may be able to extend it so llvm.can.load.speculatively could possibly return the size that is safe to load speculatively, which could then either be used within the vector loop
Any reason not to combine “can” and “load” into a single intrinsic similar to @llvm.vp.load.ff?
or peel off iterations at the beginning so that the loads in the vector loop become safe
That needs a very different semantics, doesn’t it? Or do you want to do
Peel-loop:
...
Main-loop:
assume-aligned(%lane0-ptr, vector-register-width)
if (!@llvm.can.load.speculatively(%lane0-ptr))
goto Scalar-loop; // Expect to be optimized out later based on TTI + alignment info
...
i.e. introduce extra control flow in the vectorizer with the expectation that it will be optimized out? If so, I don’t see why vectorizer itself can’t omit it based on TTI hook…
We probably should use something along these lines in llvm’s AtomicLoweringPass…
On many architectures (as recorded by the MinCmpXchgSizeInBits value in TargetLowering), there is no cmpxchg instruction with a memory access size of less than 4 bytes. So, to support 1- or 2-byte cmpxchg/atomicrmw, we round-down the address to a 4-byte aligned address, and use 4-byte memory operations. This uses a non-atomic load at the start, then an atomic cmpxchg loop.
Those 4-byte instructions may well be out-of-bounds of the object being nominally accessed. It doesn’t matter at the hardware level because the operation stays within the 4-byte granule, and nothing in the memory subsystem really deals with smaller entities. But, from a purely-IR-semantics, it’s not fully correct, since it’s doing potentially-OOB reads and writes.
From a practical POV, I don’t believe there’s any actual problem today, since the OOB write is a cmpxchg, and thus cannot modify the other memory, and since this pass runs late in the pass pipeline (after the optimization pipeline, at the beginning of the codegen pipeline). But, it may be nice if we specified some IR semantics which actually cover this usage.
How does it interact with address spaces? Is it aware of them? What if the requirements for “safe speculation” differ depending on the pointer?
Will it allow expressing “unconditionally dereferenceable” property for the loaded pointers (i.e. as if they were passed in arguments not loaded in the middle of the function)?
This looks like a nice way to make early-exit vectorization possible without relying on tricky analysis. I like the idea of separating “this load might go out of bounds” from “this is actually safe on this target”.
One small thing I’m not totally clear on: for llvm.speculative.load, is the out-of-bounds part always undef and never poison? Also, are there any concerns about SelectionDAG or later passes accidentally inferring dereferenceability again after lowering?
or peel off iterations at the beginning so that the loads in the vector loop become safe
That needs a very different semantics, doesn’t it? Or do you want to do
Peel-loop:
...
Main-loop:
assume-aligned(%lane0-ptr, vector-register-width)
if (!@llvm.can.load.speculatively(%lane0-ptr))
goto Scalar-loop; // Expect to be optimized out later based on TTI + alignment info
...
i.e. introduce extra control flow in the vectorizer with the expectation that it will be optimized out? If so, I don’t see why vectorizer itself can’t omit it based on TTI hoo
Something like that, but in that case, in the main vector loop, there would be no need to check @llvm.can.load.speculatively, as the first loop should take care of handling the iterations until the pointer becomes safe to load
That sounds like something that would be a natural extension. For the initial definition I did not yet want to pull in atomic support, to keep things a bit simpler.
Both intrinsics take pointers, so they can take address spaces into account, e.g. you could have an address space where all pointers are always dereferenceable, and @llvm.can.load.speculatively could expand to true for all pointers in that address space.
You could use a loaded pointer as address for a @llvm.speculative.load if there is some property that guarantees that the speculative load does not fault, but that is weaker than the guarantee that the pointers is dereferenceable. I am not sure it’s spelled out anywhere explicitly, but dereferenceable implies that all dereferenceable bytes are of the same underlying object.
Yes this is something where more input from people more familiar with SelectionDag would be great to double check the lowering to a SelectionDag LOAD directly is sound.
Yep, this was mostly trying to re-use the definitions of the intrinsic, to avoid having to spell out the constraints. But perhaps this can be worded generally, with `llvm.can.load.speculatively being one way to ensure the speculative load is safe?
I’m still not sure how @llvm.can.load.speculatively is any useful for peeling.
Imagine I have a pointer aligned exactly at the cache line, and I load full cache-line-size bits data from it on every iteration, but only with 0.5 x cache-line-size stride between vector iterations. It still wouldn’t allow vectorization, would it? But that knowledge isn’t encapsulated in the intrinsic itself, vectorizer would have to know that from other sources.
Yep, with the current version of the intrinsic we need to make sure that the stride matches the size, which should work well for the initial early exit use case. We can extend the intrinsics as needed in the future, I tried to keep the proposal simple & focused on the initial use case.
The intention of the intrinsic is to guarantee loads of num_bytes starting at ptr + I * num_bytes will not fault, without that it would indeed not be very useful for the vectorizer.
This intrinsic has **target-dependent** semantics. It returns ``true`` if
``num_bytes`` bytes starting at ``ptr + I * num_bytes``, for any non-negative
integer ``I`` where the computed address does not wrap around the address
space, can be loaded speculatively, even if the memory is beyond the bounds of
an allocated object. It returns ``false`` otherwise.
Legally doing out-of-bounds memory accesses is one of the most-often repeated feature requests in the category “things unsafe Rust authors would like to do but currently it’s always UB”. Great to see some movement here!
That was probably this which apparently is meant for “emulating 8/16/32 atomic operations on older ARM architectures (without atomic support) using a kernel-provided 32-bit cmpxchg function”.
Indeed, this is a problem. So what is the proposed wording for the intrinsic that takes noalias into account?
With the current wording, I would expect a llvm.speculative.load to be considered a read for all the bytes that are still in-bounds of the allocation. That unfortunately makes it unusable for most Rust use-cases since we don’t actually want to make any aliasing assertions about the “out-of-bounds” part of the load. OTOH that code always knows exactly how many bytes it is really interested in, so it is trivial to define a version of the intrinsic that makes this explicit and then only has the effects of a read for those bytes. (The Rust version of this intrinsic will almost certainly have that form.)
Even if you don’t like making this parameter explicit, I hope the intrinsic docs will say something about noalias. Otherwise we risk repeating the problems we had with undef or provenance – this kind of underspecification in an optimizing IR typically comes back to bite you later.
It should surprise nobody that I am not a fan of just ignoring problems with the specification of the operational semantics.
In particular, I would hope that the UB-aware interpreter can have support for both this operation, and noalias. (I would suggest for the noalias support to be based on this proposal.) What do you propose should it do then?
If the ambition is to verify noalias in the UB-aware interpreter[1], then speculative loads are indeed quite the complication.
I can image it is very useful to use speculative loads for something like:
; Determine if null-terminated char string %s has size 3 or less
; (and write a byte, possibly just after the string)
define i1 @is_short_str(ptr noalias align(4) %s, ptr %p) {
%v = call i32 @llvm.speculative.load.i32.p0(ptr %s)
%tmp1 = sub i32 %v, u0x01010101
%tmp2 = xor i32 %v, u0xffffffff
%tmp3 = and i32 %tmp1, %tmp2
%tmp4 = and i32 %tmp3, u0x80808080
%tmp5 = icmp ne i32 %tmp4, 0
store i8 1, ptr %p
ret i1 %tmp5
}
For the interpreter to determine if the memory access follows the noalias rules, I think it will need some oracle that can tell which bytes the speculative load actually needs to read, and which bytes are read as undef and do not affect the outcome? (In this case, the oracle could look at the memory contents of the string and return one more than the number of leading non-zero bytes, with a maximum of 4.) Or it will need to consider all possibilities simultaneously and verify that one of them does not have UB (because in practice there will be no oracle)? [This is indeed just a variation of @nikic’s example, but I wanted to illustrate that in some practical cases you do not have the %defined_size parameter beforehand, let alone that it is a constant.]
This would certainly be useful to have! Prior work in the context of restrict includes this paper (though I could not get their tool to work) and the restricted target set semantics of Homer & MacDonald. ↩︎
What I’d expect llubi specifically to do is to not treat llvm.speculative.load as a load for noalias purposes, and return the current contents of the memory at that location – which is always a valid choice regardless of which bytes are actually non-deterministic. I think this would be a sensible thing for llubi to do, as, by construction, we do not want llvm.speculative.load to ever introduce UB due to noalias interactions.
But I’m not sure whether “load.speculative.load does not count as a memory accesses for the purposes of noalias” would work as an actual specification. This effectively creates a noalias bypass for loads. noalias stores could be reordered past it, changing the result. This is actually exactly what we want (the part affected by noalias is supposed to be non-deterministic), but it’s unclear to me how such a transform would be justified by alive2.
