Semantics of affine_maps in memrefs?

MLIR
1

Hi!

I’m confused about the semantics of affine_maps in memrefs, specifically the fact that these two are valid seems to imply either the codomain of the affine_map is not uniquely determined, or there is something implicit going on:

  • memref<42x16xf32, affine_map<(i, j) -> (33 + i + 64 * j)>>
  • memref<42x16xf32> which per the documentation means memref<42x16xf32, affine_map<(i, j) -> (i, j)>>

The codomain seems to be either one-dimensional, or two-dimensional, for this memref.

As another motivating example, what are the semantics of an affine_map<(i, j) -> (i * 42, (i + j) * 3)>, where the domain of (i, j) is [0..5)x[0..5), when in a memref? What memrefs can this be attached to?

To posit the question in full generality: Let v be a value of type memref<d_1 x d_2 x ... x d_n, affine_map<f>>, with f(x_1, x_2, ..., x_k) = (e_1, e_2, ..., e_m), where each e_i is an expression of all of the x_j. What can we say about the relation between variables, and what is the semantics that we should infer about reads into this memref?

One interpretation is that there’s an implicit linearization map attached to these memrefs, that does range inference on the e_i. That is, if the range of e_i is [0 .. r_i), and we call the suffix sums s_i = \sum_{j = i + 1}^m r_j, s standing for “stride” here, then the implicit linearization map is g(y_1, y_2, ..., y_m) = \sum_{i = 1}^m y_i * s_i, and this gives us an offset into the base of the memref. This assumes that the range of each e_i is indeed computable as [0 .. r_i). The full affine map would then be the composition g \circ f, giving us a coordinate map from the range of x_1, x_2, ..., x_k, to the integers. A read into this memref would then specify k coordinates, matching the range of the x_1, x_2, ..., x_k. Is this the right understanding?

2

My understanding here is that this provides a memref view into an underlying 2-D memory space of [5*42, 10*3]. I’m not sure about the linearization in memory though, I assume this isn’t defined by the memref itself (there could be padding for example) and it may depends on the allocation.
@bondhugula is the reference on the topic!

3

You have two memrefs above. I can’t immediately tell which one you refer to by “this memref”. The map has to be one-to-one. The dimensionality of the range doesn’t have to match the domain’s, it could be 1-d, 2-d, 3-d or even higher for that shape 42x16. (Eg. tiled layouts map into a higher dimensional space.)

The part further below would answer this question. The logical space being indexed into is 5x5. The “physical” space being indexed into is at least (4 * 42 + 1) x (8*3 + 1), i.e., (169 x 25). The spec doesn’t stipulate what the allocation size should be (for eg. there could be additional for alignment), but it would have to hold at least 169x25 elements.

It can be attached to any 2-d memref since the map’s domain is two-dimensional.

This understanding is accurate. The range of the map is “physical” space and a row-major layout is assumed for it (I think this was mentioned in the doc) — that’s what the identity layout corresponds to and is defined to be. (So, an implicit linearization for the ranges is assumed.) Both, the logical and physical spaces are hyperrectangular. Also, in an MLIR affine_map, you can never divide by or take a modulo w.r.t a dimensional identifier (logical dimension id here). The upper bound for each physical dimension is always trivially computable using the lower/upper bounds of the logical dimensions AFAICS.

4

Thanks for the answer Uday. Two followup questions:

It can be attached to any 2-d memref since the map’s domain is two-dimensional.

Is it any 2-d memref, or only a 5x5 memref, given that:

the domain of (i, j) is [0..5)x[0..5)

? Is it perhaps attachable to any memref whose size is a sub-rectangle of the 5x5 rectangle? In general, will the domain of the affine_map be exactly the same as the shape of the memref it is attached to?

You also mention that:

The upper bound for each physical dimension is always trivially computable using the lower/upper bounds of the logical dimensions AFAICS.

But the lower bound is assumed to be zero? What about something like affine_map<(i, j)[s, t] -> (j - i, i - 3, i + s, j * t, -5, -5)>, where the the domain of (i, j) is again [0..5)x[0..5)? Are the bounds even statically computable as [0 .. r_i) here? The first two don’t need to start at zero, the second two depend on symbols (though perhaps that isn’t difficult if the range is only determined only after we resolve the symbols?), and the last two are just always -5. What is the computed range here?

To answer your first question, what I meant by “this memref” is memref<42x16xf32>. We can attach to it the 2-d identity map, giving us memref<42x16xf32, affine_map<(i, j) -> (i, j)>>, which per the documentation is the same thing as the original memref<42x16xf32>, or a map with a 1-dimensional codomain, giving us memref<42x16xf32, affine_map<(i, j) -> (33 + i + 64 * j)>>. This would mean that just knowing we have a memref<42x16xf32, affine_map<f>> does not uniquely determine the codomain of f. The alternative, and what seems to be actually happening, is that there’s an implicit linearization map composed after all such f, that goes from whatever the codomain of f is (and here we need range inference), to the integers. The semantics of all such affine_maps, then, after such a composition, is that this resulting integer is an offset into the base of the memref.

5

Is it any 2-d memref, or only a 5x5 memref , given that:

The map itself doesn’t have any bounds for the domain. So, an affine map with a 2-dimensional input can be applied to any 2-d shape.

Yes, of course. It’s only being applied to those integer points when attached to a memref.

But the lower bound is assumed to be zero?

All subscripts whether in the logical space or in the physical space have to be between 0 and <range - 1> (since “applying” the map to the logical space makes the physical space the logical one — this is also what the -normalize-memrefs pass does: it turns the memref into one with an identity layout and this is always possible by design.). As both the lower and upper bounds on the range are trivially computable, the map can be checked for this (when there are no symbols or dynamic shapes). When you have symbols or when the shapes are dynamic, the behavior when properties like injectivity (one-to-one’ness) of the map and non-negativity of the range are violated would be undefined behavior.