Memref not expressive enough (may need a new convention)

MLIR
1

Hello everyone,

This topic is related to the “RFC Strided MemRef Proposal”.

The TL;DR is that a MemRefType currently has no established way of tying the symbols used for dynamic sizes to the symbols used in the layout expressions.

For instance, consider the first case:
memref<?x?x?xf32, (i, j, k)[s0, s1, s2] -> (s0 + i + s1 * j + s2 * k)>
that can also be thought of as
memref<?x?x?xf32, offset:s0, strides=[s2, s1, 1]>

and the second case:
memref<?x?x?xf32, (i, j, k)[s0, s1, s2] -> (s0 * i + s1 * j + s2 * k)>
that can also be thought of as
memref<?x?x?xf32, offset:0, strides=[s2, s1, s0]>

Both forms are perfectly valid but have very different implications on the layout depending on the relationship between s0,s1,s2 and the ? in memref<?x?x?xf32>.

Assuming a memref with k symbolic sizes, I would personally be enclined to adopt the convention that the first k symbols in a layout map must be the memref sizes in their order. Optionally, additional symbols may be specified starting from s_{k+1}.

Note that this is purely a MemRefType concern and does not leak into the normalized descriptor or the ABI. As a reminder, the normalized descriptor is:

template <typename Elem, size_t Rank>
struct {
  Elem *ptr;
  int64_t offset;
  int64_t sizes[Rank];
  int64_t strides[Rank];
};

and can accomodate the different cases.

What are people’s thoughts on this?

2

Did you mean “strides” instead of “sizes” above? Because your opening statement “TL;DR …” appears to convey something different from your example. The SSA operands that bind to '?'s are separate from the ones that bind to the symbols in the layout map - but you could reuse SSA values across them. For eg:

#tiled_dynamic = (d0, d1)[s0, s1] -> (d0 floordiv s0, d1 floordiv s1, d0 mod s0, d1 mod s1)
%T = alloc(%M, %N) [%B1, %B2] : memref<?x?xf32, #tiled_dynamic>

In case you need to use a size (shape SSA value) in a layout map expression, you could do that by having that SSA shape value also appear as a symbolic operand to the alloc op. Eg:

%A = alloc(%N) [%N] : memref<? x f32, (d0) [s0] -> (d0 + s0)>
3

Hi Nicolas,

The RFC link you post seems to have problem to be accessed:

image
image2262×860 155 KB

4

@yangjunpro, I updated the link from my personal laptop. The links we get from our work accounts seem to be incompatible with the outside world…

5

@bondhugula I did mean size indeed.

This is the classical type world vs value world problem.
Notice how in the examples you gave you used the alloc op?
The claim is that the type itself does not have enough information for distinguishing between the 2 cases.
That information therefore does not pass function boundaries and analyses are subject to chains of alloc/view/subview/future_ops.

Re: " I would personally be enclined to adopt the convention that the first k symbols in a layout map must be the memref sizes in their order"

With this convention, your examples would become:

#tiled_dynamic = (d0, d1)[s0, s1, s2, s3] -> (d0 floordiv s2, d1 floordiv s3, d0 mod s2, d1 mod s3)
%T = alloc(%M, %N) [%B1, %B2] : memref<?x?xf32, #tiled_dynamic>

and

%A = alloc(%N) [%N] : memref<? x f32, (d0) [s0] -> (d0 + s0)>

respectively.

6

Is your concern that of readily knowing whether a particular symbol used in the layout map is a size itself (and thus the motivation to reserve the first few symbols for it)? Because with the current representation, you can’t determine that without looking at the defining op. If your concern is something else, then I still don’t see it.

In the two cases you were referring to: memref<?x?x?xf32, (i, j, k)[s0, s1, s2] → (s0 + i + s1 * j + s2 * k)> and memref<?x?x?xf32, (i, j, k)[s0, s1, s2] → (s0 * i + s1 * j + s2 * k)> already have different types. Dynamic sizes bound to the memref are the same kind of information as the address of the memory allocated, and so won’t be obviously reflected in the (static) type. How would adding additional symbols in the layout map help with analysis determine which SSA values were used to size a memref?

7

Yes, this is what the sentence “The TL;DR is that we currently have no established way of tying the symbols used for dynamic sizes to the symbols used in the layout expressions.” conveys (edit: updating to make it clear I am only talking about MemRefType). Note that technically there does not need to be a strict 1-1 mapping but a 1-1 mapping seems the minimal extension/convention that is useful.

Speaking of use cases, they are multiple but the one that prompted the thinking along these lines is the will to add a reshape op. Such a reshape is a noop under certain conditions and requires an alloc+copy under other conditions.

Regardless of spelling out more use cases and detailing those conditions, my claim is that this information could/should be in the type if we adopted the simple convention above.

8

I believe the initial thinking was that the sizes appearing in the layout map are just a proxy for higher dimensional maps (so if you really need them, pass them as a symbolic operand on the alloc, etc.). I would suspect that to be the case for things like reshape. I don’t think we need to discuss any more use cases - the one on reshape op you mention itself would be good enough IMO. Do you have an example on how the current representation makes it difficult to detect a no op reshape?

Whatever you propose is only adding more information to the type. I see a couple of cases where it could lead to an issue.

  1. One can no longer canonicalize the layout map to drop unused operands (because that would destroy the information you are adding in cases where the sizes aren’t used).
  2. Memrefs like these would show up as different types instead of one:
    memref<? x f32, (d0)[s0] → d0 + s0) > // size %N bound to s0
    memref<? x f32, (d0)[s0, s1] → d0 + s1) > // %N (s0) is size, %M (s1)