General support for address spaces

[tomsmeding]

LLVM has the concept of an address space on various operations that deal with memory in some fashion. In particular, the pointer type ptr has an optional address space attribute that looks like ptr addrspace(X) in the syntax, where X is either an integer or a string. In the older syntax with explicit pointee types, as in LLVM ≤14, the syntax is e.g. i32 addrspace(5)*.

The syntax with a string, e.g. addrspace("A"), refers to address spaces defined in the datalayout string; they are aliases for the appropriate numeric address space. In practice, I suspect that applications using llvm-pretty will be happy with just support for numeric address spaces.

llvm-pretty currently does not support address spaces at all. For general JIT work, explicit address spaces are typically unnecessary, but they become extremely important when generating code for GPUs, such as PTX for NVIDIA or AMDGPU for AMD, which use a variety of address spaces for the various kinds of memory on the hardware.

Would you consider adding, or accepting PRs for, address space support to llvm-pretty?

As foreshadowed in #138, we're considering using llvm-pretty for accelerate-llvm-ptx, but we are not sure yet if this will in fact happen. So no guarantees that we'll actually be using, or pushing for, this functionality in the future!

Compatibility

A priori this is of course a breaking change. For example, naively a change to support address spaces would do (among other things) the following:

 data Type' ident
   = -- ...
-  | PtrTo (Type' ident)
+  | PtrTo AddrSpace (Type' ident)
-  | PtrOpaque
+  | PtrOpaque AddrSpace
   | -- ...
+
+data AddrSpace = AddrSpace Word32
+    deriving (Data, Eq, Generic, Ord, Show, Typeable)
+
+defaultAddrSpace :: AddrSpace
+defaultAddrSpace = AddrSpace 0

eliding comments etc. for conciseness. This breaks everyone everywhere. However, the impact can be lessened by using pattern synonyms:

data Type' ident
   = -- ...
-  | PtrTo (Type' ident)
+  | PtrInTo AddrSpace (Type' ident)
-  | PtrOpaque
+  | PtrInOpaque AddrSpace
   | -- ...
+
+pattern PtrTo :: Type' ident -> Type' ident
+pattern PtrTo ty = PtrInTo (AddrSpace 0) ty
+
+pattern PtrOpaque :: Type' ident
+pattern PtrOpaque = PtrInOpaque (AddrSpace 0)

(Suggestions welcome for better names instead of PtrInTo etc.)

This way, code that produces an llvm-pretty AST will still work fine: all ASTs that it creates will have address space 0 (which is indeed the default address space in LLVM, so this is no change in behaviour), and even pattern-matching on the ASTs that it itself creates will be fine because the AddrSpace 0 pattern will always turn out to match.

Code that consumes an llvm-pretty AST will continue to work as long as the producer of the AST it ingests uses only the default address space, but if it receives an AST with non-default address spaces, it will of course crash at runtime.

All code that pattern-matches on llvm-pretty ASTs (and that is not updated to the new PtrInTo and PtrInOpaque) will of course generate warnings about non-exhaustive pattern matches, and we should probably not try to eliminate those warnings with COMPLETE pragmas.

Furthermore, of course, having such pattern synonyms is likely helpful for convenience for code that just generates simple ASTs in the default address space; such code simply doesn't have to change. :)

[kquick]

Thank you for the well thought-out proposal!

As it turns out, we have already done some work along these lines; they may not be entirely complete, but do handle some other places where address space specification is also present (GlobalAttrs, Declare, Define, and Alloca). The changes are widespread enough that they are probably addressed as a breaking change with a major version bump.

We'd propose to get this internal work realized here on a branch (in the next couple of days) that will allow us to coalesce the breaking changes, and would invite you to evaluate and contribute to that branch, with the intent of merging that branch and putting out an official release when everyone is comfortable with the updates. Does that sound like a workable proposal for your needs and timeframe?

[tomsmeding]

Cool to hear that we've been thinking along the same lines; that definitely makes things easier. :)

We're somewhat busy these weeks, but we should definitely be able to look at some code. My main concern is that we kind of know the additional expressiveness we need from llvm-pretty, but it's hard to be completely sure without doing the actual rewrite of accelerate-llvm-ptx. For example, while it is quite obvious that PtrTo needs an address space field, some noodling around with accelerate-llvm-ptx revealed that Global also needs an address space. Fortunately it is on your list (GlobalAttrs). :)

For this reason I'd be somewhat hesitant to say on a short timeframe "yes, this is all".

Other stuff we've already encountered that we will likely need:

• Parameter attributes; this is a breaking change.
• We may or may not also need attributes on a Call instruction. Usually we can ensure that the relevant attributes are also on the definition site of the function being called, which perhaps makes the call-site attributes unnecessary (?), but I'm not 100% sure we always can. This is also a breaking change.
• Some additional functionality that is less breaking:
  • Volatile loads and stores; these could perhaps be separate instructions in llvm-pretty to avoid breaking existing code
  • Some additional function attributes (convergent, inaccessiblememonly (which is curiously undocumented?))
  • Probably inline assembly in calls; I haven't investigated thoroughly, perhaps the top-level modInlineAsm is sufficient, but perhaps not.

I'd think we should handle these things in separate issues/PRs, but I wanted to mention them so that you're aware that we would need more breaking changes than just address space support. (Again, if we're really going through with this. No guarantees.)

[kquick]

Sounds good. I completely understand your inability to commit to a specific future. :)

Regarding your additional functionality: in general, we are open to improvements and would welcome contributions. I am not aware of any WIP for any of those, and your introduction on this issue is a great way to handle those: raise the issue for some discussion about proposed approaches/drawbacks, followed by any PR's as you discussed. The "Discussions" area is also a great place to develop initial concepts into actual plans.

My general hunch is that while there are several breaking changes you've identified, we may want to roll those out in incremental major releases (where it makes sense), allowing library consumers to perform incremental upgrades on their end instead of a broad set of all-at-once changes that would require more comprehensive modifications for consumers (although they are free to "skip to the end" if they prefer that approach).

Let's stay in touch!