RFC: Changing byte-alignment to bit-alignment for memref and vector dialects

MLIR
1

Motivation

This RFC proposes extending the alignment attribute to data sizes smaller than a byte. Types and data smaller than a byte are already being used in several dialects (e.g., vector, amdgpu, rocdl, spirv), in passes to emulate these types (–test-emulate-narrow-int), and being used in several targets. To support alignment for data sizes smaller than a byte, we propose changing the byte-alignment attribute to be a bit-alignment attribute.

Byte Alignment

We say that some data is n-byte aligned when it is stored at an address a and a is a multiple of n where n is also a power of two.

Some targets require data of certain size to be n-byte aligned. When loading or storing data of these sizes from or to addresses which are not a multiple of n, these targets may result in a performance penalty or an alignment fault. The alignment attribute is useful to allow compilers to properly meet alignment requirements to optimize and generate valid programs.

Bit Alignment

All byte-alignment attributes can already be expressed in terms of bit-alignment. The conversion of a byte-alignment attribute to a bit-alignment attribute is just the byte-alignment attribute multiplied by eight.

Element Alignment

An alternative to bit-alignment that would allow for sub-byte data to be annotated with an alignment attribute is to have the alignment be defined in units of the number of elements.

E.g., For a vector <N x i4> with element-align 2, then the bit alignment is 8, and so the byte alignment is 1.

Table of Pros and Cons

Byte Alignment Bit Alignment Element Alignment
Pros Status Quo is preserved Optimized lowerings for vector dialect. Easy conversion between byte and bit alignment. Optimized lowerings for vector dialect.
Cons No optimized lowerings for vector dialect Not so easy conversion between byte and element alignment.

Next Steps

2

cc: @krzysz00 @ftynse @banach-space @dcaballe

3

This is interesting. I’m intrigued by the consequences of having non-addressable objects. Can you elaborate on that? How would we lower memory accesses to them?

4

Thanks for sharing this!

This part is a bit unclear to me. What specific optimizations are we currently missing, and which ones would become possible (or easier) with bit-level alignment? Could you elaborate a bit more on this?

On a related note:

Since most of my work lowers to LLVM, the current byte-based alignment design aligns well with my needs. I’m guessing your use case might not involve lowering to LLVM directly? That would make sense, I’m just trying to better understand your context and explain where I’m coming from.

-Andrzej

5

To add to the usecase for sub-byte alignment, the narrow type emulation pass - which deals in things like loading from a memref<... x i4> by rewriting all such memrefs to be i8 memrefs and rewriting all the operations - has to be pessimistic with, for example, vector.store %v, %m[%idxs...] : vector<2 x i4>, memref<... x i4>, because it can’t know whether or not %idxs... gives a byte-aligned value or not. If that store were known to have alignment = 8 bits, you can just rewrite it to a store of a vector<1 x i8>, but if it has alignment = 4 bits, you have to do a bunch of atomics on two adjacent bytes.

The problem with defining alignment in terms of bytes is that there’s no way to represent that 4-bit-aligned case.

I think having the alignment be in terms of bits is a good solution that allows us to express memory that contains sub-byte types, which often occurs in ML contexts. When lowering to LLVM, we’ll probably want to follow LLVM’s conventions for non-power-of-two loads (which, IIRC, round up to the next multiple of 8 for the purposes of calculating a natural alignment) … or we could just throw an error or a warning because you probably didn’t mean to do that.

6

Sorry, I still can’t follow… if alignment = 4 bits leads to inefficient vector code, same as not providing any alignment information, what’s the value of adding alignment = 4 bits?

7

As an engineer working for AI accelerator, there can be cases that we need to align for bit smaller than 8 bits. Although LLVM assumes data load & store is always done in byte level, and most of general-purpose hardware (Such as CPU or GPU) do so, for hardwares optimized for efficiency, load & store with size less than 1byte, or aligning them with less then 8 bits may be beneficial for writing compilers for them.
Typically, for running LLMs, many people are interested in quantizing them less than 8 bits, and their weight can be stored in 4 bits which is supposed AI accelerators.
Yes, storing vectors in 4bit alignment with CPUs may result in efficient code, but for custom-designed hardwares, it can be another story.

8

The thing is, alignment = 4 bits is not equivalent to leaving the alignment unspecified. Leaving the alignment unspecified means you take whatever the default alignment is. For example, on the way to LLVM, the natural alignment of an i4 is 1 byte or 8 bits, and there’s no reason that MLIR’s operations couldn’t round up for sub-byte types by default.

(That is to say, on the way to LLVM, an operation that’s explicitly aligned to less than a byte probably should be an error because you should’ve rewritten that into an actually performable operation before you hit that point … but an operation with unspecified alignment shouldn’t trigger that behavior.)

9

Generally, I think it is okay to change to bit alignment on memrefs, the complexity of lowerings doesn’t seem to increase that much and it supports targeting not only LLVM.

One thing I’d like to understand better is whether this affects or solves the alignment quirks inherited from LLVM IR of vectors with non-power-of-two sizes, e.g. memref<?x vector<3xi8>> that at some point was expecting the vector to be 4-byte aligned at every index or, funnier memref<?x vector<6xi4> and memref<?x vector<4xi6> that IIRC were 8-byte and 4-byte respectively.

Another random memory is that some folks were adamant on non-eight bit bytes back when I was working on data layouts. I don’t know of any MLIR user that actually needs this, but maybe they exist…

10

I’d like to understand this a bit better, this is probably worth an ODM with a presentation showing the motivating cases and the various options moving forward. Is someone up for this?

I have to say I’m a bit confused about sub-bype alignment, this seems to assume that a pointer address granularity is smaller than a byte in the first place, right? I haven’t worked on such platform, are there public example of such platforms? Can C/C++ even support this?

11

Answered in reverse order

  1. I don’t think there’s anything in C/C++ that stops you from having a bit addressable-architecture where sizeof(char) == 8?
  2. I’ve got the vague sense you can get really clever with addressing on FPGAs and such?
  3. So, more fundamentally, re bit-addressability … even if the underlying architecture or backend doesn’t let you address less than a byte, it is very useful at the application level to be able to act as if you can address less than a byte and then to rewrite that into accesses to parts of (a) byte(s) … a tensor<64x32xi4> is a perfectly sensible thing to pass to a ML kernel these days even if you don’t have a 4-bit load. The ability to discuss alignment as if you could load/store less than a byte lets you implement the common case where you’re, for example, dealing with byte-aligned or 32-bit aligned vector<8xi4> or the like efficiently, while still giving us a way to express “yeah, my code might want to just change the high half of a byte, please emulate that”.

I think this will affect that, in that it’ll let you more easily write the MLIR passes that translate such a memref into memref<... x i8> or such, thus allowing a packed representation of this sort of vector. If you do lower a “weird” memref like that to LLVM … you’ll get the LLVM assumptions, but also that’s a “it shouldn’t work, don’t do that” situation.

12

What about these excerpt from the C++ standard?

The fundamental storage unit in the C++ memory model is the byte.

>The memory available to a C++ program consists of one or more sequences of contiguous bytes. Every byte has a unique address.

Also the C standard on “memory location” is pretty explicit on the definition of a memory location:

either an object of scalar type, or a maximal sequence of adjacent bit-fields all having nonzero width

The fact that there are no scalar type smaller than a char (8 bits) seems to invalidate any sub-byte addressing? (all the bitfield examples in the section seem relevant to me as well).

I understand that this is the claim behind this RFC, but the request is to elaborate on this with motivating example actually. For example bitfields in C look like you can address and load individual sub-byte elements in the language, but the LLVM optimizers is oblivious to this (the frontend has to make it byte-aligned). The claim here is that somehow what we do in MLIR with memref would benefit from a different model than LLVM: it’s very possible but some more explanation seem warranted to me.

Beyond the motivating example, I’m wondering about lowering things like:

for %idx from … to …:
  %value = load %ptr[%idx] : memref<?xi3>
13

Ok, so (even though sub-byte hasn’t been rigged for non-powers-of-2), that load can be implemented as follows - and, if i3 were i4, we do have MLIR passes for this today

for %idx from ... to ...:
  %idx.byte = (%idx * 3) /[floor] 8;
  %value.byte1 = %ptr.byte[%idx.byte] : memref<?xi8>
  %value = if (%idx mod 3 == 0) {
    yield trunc %value.byte1 to i3
  } else if (%idx mod 3 == 1) {
    yield trunc (%value.byte1 >> 3) to i3
  } else if (%idx mod 3 == 2) {
    %word1 = vector.load %ptr.byte[%idx.byte] : vector<2xi8>, memref<? x i8>
   %word1.scal = bitcast word1 to i16
   yield trunc (%word1.scal >> 6) to i3
  }
  ...
}

The bigger issue is actually if that loop were a store, where you need CAS loops to ensure you don’t lose modifications to other byte parts.

However, if we modify your example to

for %idx from ... to ... by 8 {
  %val = vect.load %ptr[%idx] : vector<8xi3>, memref<?xi3>
  ...
}

In general, we don’t know that %idx is aligned to a byte boundary. We have do some complex handling of the different cases of %idx mod 3 to implement this.

However, if I were to have

%value = vector.load {alignment = 24 bits} %ptr[%idx] : vector<8xi3>, memref<? x i3>

, then that lowers to

%value.byte = vector.load {alignment = 24 bits} %ptr.byte[(%idx * 3) / 8] : vector<3xi8>, memref<? x i8>
%value = bitcast <3 x i8> %value.byte to vector<8xi3>

And, yes, on the way to LLVM, alignment = 24 bits becomes align(1).

And for a practical application of all this … f6E3M2FN and all the other 6-bit float types are basically this example - you’ll always want to load at least 8 - though usually 32 - such 6-bit values at a time, and this RFC gives us a way to indicate whether or not you use a slow or fast lowering for such a load.

14

That’s assuming that the pointer is byte aligned in the first place isn’t it?

Anyway: I wasn’t asking about the feasability of lowering it: I’m concerned about the abstraction hiding a very costly lowering here, and question the suitability of this abstraction.
I can’t really make a tradeoff analysis as I’m still waiting for the motivating example :slight_smile:

for %idx from ... to ... by 8 {
  %value = vector.load {alignment = 24 bits} %ptr[%idx] : vector<8xi3>, memref<? x i3>
  ...
}

This example here is one where the alignment is still a multiple of 8 right? For this proposal, we’re looking for an example where it would not be.

Side-note / question, I can’t find in our documentation the expected memory layout of memref<?xT>? I assumed that every T in the memref should be aligned to the natural alignment of T (same rules as a C array basically). But that means at minima having a byte-storage for every T, which I assume you’re trying to not do for these sub-byte types.
It’s hard to reason about this because as mentioned before, I don’t believe C/C++ have such concepts and so we don’t have much point of comparisons.

15

I think the memory layout of a memref<?xi4> is somewhere between “we never said”, “it’ll depend on how you lower it”, and “it’s the obvious packed representation” (by way of EmulateNarrowTypes doing bitcasts to tensors of bytes for memrefs of i4, and thus that being a pass that is implicitly making claims about what a memref of i4 is … if you use that lowering or a similar one).

One reason I want bit-alignment is so that

  1. Lowerings that see alignment = 3 bits on a load of an i3 or a vector<8xi3> will have that as a signal that they have to go down the inefficient path
  2. Therefore, if I’m doing some reasoning to determine the alignment of each load/store, I can tell something’s gone wrong if I actually end up with a bit alignment

Basically, the case where you have to do an inefficient lowering is unrepresentable in MLIR today, even though it should be in order to allow the code that dispatches between the fast and slow cases to work correctly.

Loading a single i4 or an i3 from a densely-packed memref of such things is an operation that conceptually makes sense, but has wildly differing implementations depending on what you know about the bit-level alignment of the access.

Or, to your natural alignment point … the natural alignment of a 4-bit value is dependent on where in the compilation you are and what you’re targetting. At the level of your ML kernel compiler, the natural alignment of a 4-bit value is 4 bits, because you know you’re passing in densely packed buffers. At the LLVM level, the natural alignment value of a 4-bit value is 1 byte, because that’s the best you can do because memory is byte-addressable. To get between those two worlds, you no do rewrites.

(I’m also extremely fine with actually making the statement that a memref<? x T> is a densely packed array of T with no padding, thus solving all of this and therefore making memref<?xi4> on the way to LLVM a lowering error)

16

I think it may be reasonable to enforce alignment during lowering. For sub 8-bit lowering cases, how they will be aligned looks like “black box” for frontend developers. In other words, it’s unclear how it will be lowered, which means it will result in unexpected consequences depending on the target.

However, if we think of this as “adding requirement for lowering pipeline”, there are several similar examples. For example, lowering unsupported floating point types (such as f16 for hardware only with bf16) will result in frequent bitcasts, or call to compiler-rt. Frontend developer might not expect this for optimization, but compiler does it automatically.

Therefore I think this question is related to “If compiler is doing something with potential side-effects, is it a problem? or a feature?”. For your 3bit alignment example, lowering vector<8xi3> in inefficient way may be feature of the compiler for increasing flexibility, but we might say this is an error because compiler should not sacrifice the performance.

Your suggestion potentially gives people way to determine whether compiler will treat this an error if it has to do different alignment other than 3 bits depending on the hardware. In that sense, this may be useful.

17

I’m uneasy with “it’s the obvious packed representation” because it’s inconsistent with how we are used to work with any array in memory with C/C++, where elements are aligned and padding is assumed in between elements.

That is: we could also just say that the only valid “packed” representation is memref<?xvector<2xi4>> (or any multiple of 2 for the vector size).

The alternative is “we don’t support it and people need to work with 8bits alignment when they produce the IR in the first place” (which is consistent with C/C++ programming model which don’t have subbyte-aligned pointers AFAIK).

Slight twist: it is not doable with the type/operation that are available upstream right now, one can represent it in MLIR.
Whether it needs to be supported by all the upstream memref/vector dialect is the open question here: we can also push the complexity out of the system in the lowering to these dialects (which would keep things in terms of actual “addressable objects”, compatible with C/C++).

I don’t see any alignment lower than 1 byte as “a thing” actually (other than a conceptual high-level mathematical thing). So to me the natural alignment of a 4-bits value is 1 byte, and packed storage produces non-addressable values: not “4 bits aligned individually addressable value”.
Just like bool is handled in C++: you can pack it, but there is no such thing as a “natural alignment of 1 bit for booleans” (this is somehow the whole std::vector<bool> abstraction question).

That seems like a problem to me: especially since we’re talking about interacting with memory buffers, we need very clear expectations of the layouts, this can’t be something opaque and internal to the MLIR compiler (on the contrary of things like whether a local variable is on the stack, in a register, or promoted to the heap even, which are non-observable to the user).

18

In ML workloads, we want to have the notion of a tensor<...x i4>, for example, and those are densely packed. Accessing an arbitrary index within a tensor is a well-defined operation - it’ll involve some bit-manipulation in the general case. Those lower to memref<... x i4>, and we have - we currently have - passes that will rewrite all operations on a memref<... x i4> to those on an equivalent memref<... x i8> in order to get us back into the C/C+±like world where memory is made of bytes.

Before that, memref is still a somewhat abstract concept, and because memory that consists of 4- or 6-bit values or such, densely packed, does exist and is used, we have a representation for it - memref<... x iN>. The notion of bit-alignment allows us to work with such abstractions more efficiently.

(And yes, I agree that memref<... x i1> is ugly here, because there are two possible interpretations it could have - it could be a memref of bits or a memref of bytes that are all 0 or 1. Looking around, we’ve currently got an ad-hoc exception to the notion that everything’s densely packed, but I also have a vague memory of downstreams with rewrites that try to deal in memref-of-i1-is-bits)

So, yes, there’s some slipperiness around what exactly something like memref<... x i4> or memref<... x i3> means in terms of memory layout (or, really, memref<... x i1>).

I don’t think we need to solve that today - or ever. Natural alignment’s already a target-dependent thing (see LLVM having data layout keys for it) and so, because memref in an abstraction, we can just continue down the current course of letting the exact details of how the values are arranged in a memref be lowering/target/application-dependent. And since one possible and useful usage is densely packed sub-byte qualities, allowing alignment to be expressed in bits makes such a usage more useful without harming the usability of other possible modes (like the padded version).

That is to say, strong reject of defining memref<... x i4> as padded - it’s too useful to get rid of.

(Edit: to put it another way, the MLIR standard library is a toolkit for building compilers. There are several valid data layouts for a memref<... x i4>, and the person building the compiler is responsible for picking. The standard library should work with any of the reasonable choices - and, between most passes just transparently not minding a sub-byte-value memref and EmulateNarrowTypes existing - it already does. So I don’t see why we’re suddenly diverging into standardizing the memory layout question, since it has never been fully defined for all downstream users and still doesn’t need to be with this proposal.)

19

You seem to be claiming that this is an intentional thing in the design and should be supported, I’m just not convinced yet because that seems like a thin argument.
To me, it’s not clear that a memref<... x iN> should be densely packed. Someone adding code making this assumption isn’t a design decision here!

How so? I saw your example of lowering and it seemed all but efficient to me. What can be done more efficiently by allowing this?

Memref started as an abstraction over memory, that is pretty fundamental to me that we can address the objects inside a memref: now we never discussed it before, some people started (downstream) to do ad-hoc thing with this, but that does not mean it is suddenly becoming “blessed” in the design.
So the question seems completely relevant to me to this RFC as this RFC makes a premise that it should be supported.

20

I claim that memref<... x iN> being densely packed is a permitted semantics for memref<... x iN>, not a required one. The primitives available upstream do support (to a limited extent, via rewrites) that interpretation.

Furthermore, such semantics are a useful interpretation of memref that people are using and need to use.

This RFC is merely extending the notion of annotating the alignment on memory accesses to better-support this existing usecase.

I realize that “the exact layout of a memref<... x i{N < 8}> or memref<... x f4E2M1FN> is unspecified upstream, use the lowerings and rewrites that match what you and the user of your program agree on” isn’t a particularly satisfying answer from the point of view of language semantics or being to able to do really firm reasoning about your IR, but there comes a point where such questions get in the way of building something useful. The ability to talk about something that’s logicially interpretable as a dense array of 4-bit integers is useful, even if you can’t output C or LLVM that operates on such an object - although, since the version of the world where there’s padding also makes sense, we shouldn’t disallow it.

I’m arguing that, because there are several viable definitions for the layout of a memref<... x i4> , we shouldn’t pick one as the one true correct answer, and stick a note onto mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp that says that this only works if you’re using densely-packed sub-byte values.

That being said, the really principled thing to do here would be to lean on DLTI a bit more, and actually make people specify whether they’re a densely-packed-i4 or a padded-i4 module/function/… - and that’d resolve the ambiguities around i1 as well. That is, if we wanted to formalize the target-/context-dependence of memref layout, we’d ask people to specify align(i4) = 4 bits or align(i4) = 8 bits in a data layout or the like. However, in the interim, improving the expressibility of code that works with the packed case doesn’t seem like an issue to me.

(Heck, this packing vs padding issue can even come up outside of sub-byte types. If I were targeting a word-oriented architecture, whether align(i8) = 4 bytes or align(i8) = 1 byte is also a decision I need context to make, and memref will behave no matter which choice I made. memref stands silent about the exact means of computing the address of memref.load %m[%x] - you just need your callers to agree with whatever that lowers to.)