Grab-bag of improvements to the IR we send to LLVM

[krzysz00]

Notes to reviewers:

• The first four commits (related to DistinctAttr and the handling of AliasScopeAttr and the like) are direct ports from upstream: I figured I'd port in now because the new solution is nicer and it'll save us the pain during the next upstream merge.
• The three commits after that (related to adding nontemporal to vector.load and vector.store, improving the -gpu-to-llvm attribute handling, and adding !invariant.load to llvm.load) are mine and are upstreaming candidates that should be reviewed as such.
• The last commit is all the internal plumbing, both in terms of -rock-analyze-memory-use and -rock-prepare-llvm`. I've copied out the commit message for that last commit below.

---

A lot of our existing annotation scheme for memory accesses was
scattered all over the place. This commit unifies things into
-rock-analyze-memory-use.

This pass covers:

1. Walking through global_load and global_store operations in order to
   find all the "needs 64 bit index" attributes and moving that
   declaration to the function level. That means we won't have issues
   with the fact that that runs in blockwise-gemm-to-threadwise.
2. Detecting readonly and writeonly arguments. This is done by looking
   for arguments that are only (ignoring views) arguments to global_load
   or global_store(set) respectively, and then putting those annotations
   onto the function arguments for lowering to LLVM.
   2b. Setting nontemporal flags on stores to writeonly values. If we
   won't be reading it, we want the write to skip cache, and nontemporal
   is a hint to make that happen. (It doesn't apply to buffer stores yet
   since we don't have address space 7, but it's better than nothing).
3. Setting all sorts of relevant LLVM annotations (noalias, nonnull,
   noundef, dereferencable, and so on) in order to promise the compiler
   things we know are true anyway (if they're not, that's on the client).

Then, this commit also expands our LLVM preparation pass.
The pass is made to run on llvm.func ops, and we use the nesting
mechanism to ensure that we run that on functions inside GPU modules
specifically.

The preparation pass is updated to

1. Update the alignment of vector loads and stores to be the alignment
   of the vector, not its component, since MLIR's lowering is
   pessimistic in a way we don't want to be.
2. Set !invariant.load on loads from readonly values to given out even
   stronger hints about the ability to hoist or sink operations.
3. Work around a shortcoming in how the AMDGPU backend handles
   noalias (specifically, it gets dropped on kernel arguments during
   target-specific lowering) by adding our own alias scopes, which
   indicate that loads from one pointer or buffer can't alias the loads
   from the other ones. This will hopefully reduce the number of
   scheduling DAG edges and therefore give us some of our compile
   performance back ... or at least it'll give us better codegen thanks
   to more accurate alias information.

This commit also adds the "nontemporal" attribute to rock.global_store
so that -rock-sugar-to-loops can set it on the underlying stores.

[krzysz00]

Also, @giuseros, if you've got the command for benchmarking that post-large-tensors slowdown handy, I can run it on this change.

[krzysz00]

Update: preliminary perf checks are showing that setting nontemporal is usually, but not uniformly, a bad idea in our contexts, so that part of the PR will probably be removed (though I'll keep the upstream changes)

[giuseros]

The code looks fine, I am a bit worried about introducing further analysis in the back-end that might slow compilation time. In particular, I am worried about the invariant.load attribute (see rust-lang/rust#103070). Could you try to tune some configurations and verify that the compile time does not get too bad?

[giuseros]

Not sure how much is relevant, but I stumbled upon this: rust-lang/rust#103070 , which talks about invariant.start. Do you mind making sure that the compilation time is not heavily affected by this?

[giuseros]

See my comment above :)

[giuseros]

Also, @giuseros, if you've got the command for benchmarking that post-large-tensors slowdown handy, I can run it on this change.

This is the command I was using to investigate compilation slow downs:

time python3 ./bin/tuningRunner.py --op gemm --config="-g 64 -m 1024 -k 1024 -n 384 -t f32 -out_datatype f32 -transA 0 -transB 0"

[manupak]

It d be much easier to review if you can break the things you port from upstream and what you've implemented new in this PR, because the former would be a no brainer to get in.

[krzysz00]

@manupak You can set the review view to be filtered by commits, if that helps

[giuseros]

LGTM!