Drop the requirement of having to return `AffineMap`s from `MemRefLayoutInterface`

Posting here to signal a change that I wanted to make with Allow MemRefLayoutInterface to return null AffineMaps by MaheshRavishankar · Pull Request #162757 · llvm/llvm-project · GitHub . The current state is that all Attributes that implement the MemRefLayoutInterface need to return an AffineMap that represents the layout of MemRef. While layouts that can be represented as affine_maps are an important category, there can be layouts (not necessarily in-tree) that might use layouts that are not representable as affine maps. I wanted to change the requirement to be that Attributes that implement this interface can return an AffineMap if possible, but returns a null value when the layout cannot be expressed using affine maps. I am not suggesting any fundamental change to memref semantics (there was a long discussion on this previously [RFC] MemRefType affine maps - list vs single item ) . All that changes is that transformations that rely on the AffineMap representation of layouts just need to handle the failure gracefully. I didnt find many cases where things need to be fixed in MLIR (all tests have passed on the change above), neither did I find anything was needed in IREE ([Do not submit] Checking llvm/llvm-project#162757 by MaheshRavishankar · Pull Request #22269 · iree-org/iree · GitHub) suggesting that this change can be made without much downstream impact.

cc @ftynse

1 Like

I think this is in line with the progressive change of the memref layout to be less bound to affine, especially after having the strided layout be the default.

How so? As far as I remember, strides layouts were considered when designing this interface: the idea was that any kind of layout attributes was possible, but with a restriction that it should be convertible to an affine map.
I don’t see this change as trivial right now, and I would like to see more motivation for it, in particular in-tree.

We moved from memref having the layout as a chain of affine maps (with some affine-specific composition), to a single map, to having a strided layout encoded as an affine map that we constantly had to “parse”, to having a strided layout as a first-class attribute though still convertible into an affine map. Relaxing the convertibility requirement is a natural next step in this evolution. Especially given that upstream transformations predominantly rely on the strided layout, whereas affine-specific layout seems to be mostly used in (non-public) downstreams.

I followed until here, I don’t quite get this step, can you elaborate why this would be a natural evolution?
I am also still interested in a more clear motivation as well: the fact that we could do it does not mean we should, your “natural next step” seems to fall a bit short on this aspect right now.
The kind of things I would expect here is around the following non-exhaustive kind of questions (basically my comment is that this should be a more fleshed out RFC IMO):

  • can we have some example of layout of interest that can’t be converted,
  • how would these layouts mesh with all the memref infra and assumptions that may exist with the implicit affine maps conversion (that can be non-explicit sometimes, but embedded in the mental model formed around memref manipulation as well)
  • do we have plans to have such layouts upstream,
  • etc.
1 Like

The entire direction so far has been carefully detaching memref from affine structures that is has been carrying since inception (as a polyhedral compiler pointer). Not mandating convertibility to affine maps is perfectly in line with that.

I am not the author of this RFC.

However, I expect this change to be similar in scope and style as [RFC] Memref of custom types .

am also still interested in a more clear motivation as well: the fact that we could do it does not mean we should, your “natural next step” seems to fall a bit short on this aspect right now.

I consider it mostly legacy to say that the strides should be convertible to AffineMap. With the introduction of the interface itself it opened the door to not using AffineMap as a way to represent strides. This is going a step further and saying it is optional that the strides be representable as AffineMap . Utilities that require the strides be representable as AffineMap just have to handle the case where it isnt representable this way gracefully.

There are many such examples. The underlying buffer could have an opaque layout, it could have a data-dependent layout. At this time I am experimenting with such layout options, and need a bit more experimentation to have a more coherent story. Irrespective of that this restriction of “must be convertible to AffineMap limits the ability to even experiment. Happy to discuss those in due course, but those IMO need more proof/require to be worked through more. Before that though, as Alex said, this is one step towards :

Having convertible to AffineMap can provide you additional benefits, but does not need to be a requirement.

Representability/convertibility to an affine map was motivated by the fact that all layouts of possible interest at that time, including complex ones not immediately of interest, could be represented as affine maps. So they are very general and powerful, given that the affine maps in MLIR allow mod’s, div’s, and even semi-affine expressions, while providing a way to de-abstract to generate subscripts into the physical space mathematically.

Obviously, affine maps won’t encompass all layout functions, and if there is a need to represent something even more general (e.g., functions that aren’t even in closed form, layout mappings stored in global constant arrays, etc.), removing the restriction may make sense. But I don’t see it as a natural evolution, but motivated by things where the current infrastructure is inadequate, like Mehdi says. So I have the same exact questions in Mehdi’s bullets.

The choice of affine layout maps isn’t really due to “polyhedral compiler” considerations in the current state, and for a long time. In fact, the layouts could be (and are) used without having anything to do with notions of polyhedral analysis/transforms. (In fact, semi-affine is not even affine and analyzable to the latter’s extent.) It’s the availability of affine expressions in the IR as a first-class thing, coupled with their generality to represent a very broad class of layout functions (described in the first para above), that remains as the lasting motivation.

Do you really need to completely remove the restriction on the memref layout to experiment? E.g %A[%L[%i]] could be viewed as the result of a data-dependent layout (sitting in%L) applied to %A in line with your example. One could materialize an access with such a layout without changing the current memref layout’s structure/restrictions. But it would require realizing the layout instead of carrying it in the type system. I don’t see the point about “opaque layout” though. At the time you have to lower, generate code, realize the layout, etc., you will need to know what it has.

The choice really appears to be between: carrying around arbitrary things as a layout in the type system, or just deabstracting/materializing the layout right away (into the memref uses) as soon as you know it’s what you want. In the current state, the memref type carries MemRefLayoutInterface-implementing Attributes as the permitted layouts. Also, if we want to really generalize that, it would be along “Attributes that implement the MemRefLayoutInterface need to return an AffineMap or something specific/more general that represents the layout” instead of “AffineMap or null” – because the latter (returning null) is completely leaving semantics unspecified and raises the second bullet question that Mehdi lists in comment #5.

Not sure how you would do that.

Would it work if I just add an attribute that is a UnitAttr that says the layout is opaque instead of leaving it NULL? I was using NULL as a representation of that.

I meant %s = memref.load %L[%i] and memref.load %A[%s] - i.e., the layout is “materialized”.

But that’s exactly what I said would be the issue. Saying the layout is opaque doesn’t specify what the layout is. Separate from the unit attr, where exactly are you specifying what the layout is? i.e., where in “physical” memory a given logical indexing maps to with that opaque layout. Today, the MemrefLayoutInterface is fully specifying it.