LLVM `argmemonly` attribute

From the LLVM docs:

argmemonly

This attribute indicates that the only memory accesses inside function are loads and stores from objects pointed to by its pointer-typed arguments, with arbitrary offsets. Or in other words, all memory operations in the function can refer to memory only using pointers based on its function arguments. Note that argmemonly can be used together with readonly attribute in order to specify that function reads only from its arguments.

Could we assign this attribute to certain functions in rust to aid optimization? I don't think it gets added automatically during LLVM's optimization. I suggest we could ascribe the attribute when the following conditions hold:

  • Function does not transitively call any unsafe functions or blocks
  • Function's arguments (as they are called by the abi) have a "reference level"* of at most 1
    • I would need to clarify whether e.g. a function taking &&u8 can read the u8 or not. This is not 100% clear from the docs
  • Might need to worry about (immutable) statics?

* The "reference level" of a type is defined as follows:

  • Scalar types have a reference level of 0
  • &T, &mut T have a reference level of 1 + reference_level(T)
  • *const T, *mut T have a reference level of 0
  • A struct or enum has a reference level equal to the maximum of its elements'.

What do you guys think about this? A path worth persuing?

2 Likes

The conditions you enumerate are not sufficient. The two main omissions I can spot right now are:

  1. For foo to be argmemonly, all functions it transitively calls also have to be argmemonly.
  2. Besides statics (mutable or not), many constants are translated to global memory. So a function that uses, for example, a string literal, is (likely) not argmemonly because it accesses the memory of the string literal. (See @byte_str.0 in the debug LLVM IR.)

If you load a pointer through a pointer-to-pointer, the loaded pointer is not based on the pointer-to-pointer, so &&u8 would be disqualifiying.

Is this necessarily true? Maybe the condition is a lot stronger than I realised - but I would expect the following to be argmemonly:

fn not_argmemonly(x: &&u8) -> u8 {
    **x
}

fn is_argmemonly(x: &u8) -> u8 {
    not_argmemonly(&x)
}

If this is not right then it's a lot less useful than I thought.

This makes sense. It's curious that inlining constants into a function might allow better optimizations. To me it seems the same thing to say let x = b"string"[0] and let x = b's' and one should not be penalised over the other.

Thanks for clarifying. So "reference level" is required I guess

You're right, the second function is argmemonly. I jumped ahead to trying to create a sound rule, but it's not an exact one. The actual underlying issue I was getting at is this: You need to ensure that other functions which are called transitively also only access memory passed as argument to the function you want to mark as argmemonly. For example, if foo calls bar and bar accesses global memory, foo can't be argmemonly.

The reason for that is: Generally, a property like "the function does not do X" means that "during the execution of the function, X does not happen". Put differently, it doesn't matter whether X happens literally in the function itself or is delegated to another function — the effect from the caller's point of view is the same. That's because the caller's point of view is what matters if you want to optimize a call to the function that is claimed to not do X, and also because function calls may be inlined.

Well, after you've optimized the former to the latter, they are equivalent. The distinction only exists if you analyze the source program without any simplification. One optimization allowing further optimizations that were hampered by the code that was optimized away is a common theme.

Ah, right - so it can't transitively depend on anything which reads a global static/constant, nor executes (particular types of) unsafe code. This is more restrictive than I'd hoped.

I suppose what I'd really like is a slight weakening of argmemonly that allows read-only access to global immutable constants. I suspect that a lot of optimizations relying on argmemonly would still be valid under this weakening.

LLVM’s “argmemonly” attribute is actually ok with constant memory. What it really cares about are memory dependencies (load-store, store-load, store-store), and since there can never be a store to constant memory, loads from it never have dependencies.

So in your example, above, the LLVM IR has

@str.0 = internal constant [3 x i8] c"foo"

which is constant, so it’s ok to reference from an argmemonly function.

3 Likes

This is great if true. It means (up to some things I'm sure I've missed) that the first 2 points are sufficient to mark a function as argmemonly. It makes sense from the point of view of optimizations for it to be this way, but it's not technically what the docs say so I'd want an official word on whether this is allowed or not.

I want to start playing around with this. Does anybody have any pointers how I would start evaluating whether a function is "transitively unsafe-free"? Is this information encoded in MIR or is it lost by that point? How would I pass this along as an attribute?

This topic was automatically closed 90 days after the last reply. New replies are no longer allowed.