Linear transfers (without TensorAccessPattern) shouldn't use data layout transformation dimensions

[andrej]

Documenting this because I just ran into this:

If we have a large linear transfer of a type that is described as a multi-dimensional tensor, e.g. (M, K), when using iron.fill, the lowered code ends up using the data layout transform dimensions:

    ty = np.ndarray[(M, K,), dtype_in]
    with rt.sequence(ty) as (inp):
            rt.fill(fifo.prod(), B)

produces:

        aie.dma_bd(%arg1 : memref<32x2048xbf16>, 0, 65536, [<size = 1, stride = 0>, <size = 1, stride = 0>, <size = 32, stride = 2048>, <size = 2048, stride = 1>]) {burst_length = 0 : i32}
        // (note that the two lowest data layout transformation dimensions are used.)

When one of the dimensions is too large (e.g., K > 1024), this leads to this error:

error: "-":29:9: 'aie.dma_bd' op Size 0 exceeds the [0:1023] range.

The hardware can be used in a "layout transformation" and "linear" mode. (For relevant special handling of linear transfers, grep for isLinearTransfer in the codebase.) In linear mode (I'm just calling it that, not sure it has an official name), the data layout transform dimensions are ignored and we can thus do larger, longer transfers without running into above error.

Currently, there are two workarounds:

1. Change the type to a single linear buffer, in the example above ty = np.ndarray[(M*K,)], dtype_in] rather than ty = np.ndarray[(M, K,), dtype_in].
2. Use a simple linear TensorAccessPattern with the fill: B_tap = TensorAccessPattern(tensor_dims=(M,K,), offset=0, sizes=[1, 1, 1, num_batches * K], strides=[0, 0, 0, 1])

I think it would be a usability improvement if for simple linear transfers (without any TAPs given), the second workaround was implemented as the default lowering.

cc @hunhoffe

[hunhoffe]

Nice documentation of this issue! I like the idea of the second option as a default lowering. I do think there are other opportunities -- not just for linear transfers -- for sizes/strides to be modified from a version disallowed by the hardware to a version allowed by the hardware but that describes an equivalent transformation. If I remember correctly, there is more progress towards this in air? It would be neat to consider this issue more deeply and see if we could borrow some of air's transformation logic to solve this problem for a wider variety of situations, but the one you described seems common enough that it might make sense to special case it to get a more immediate solution.

[andrej]

Splitting up dimensions when we hit limitations would be an additional useful feature. But [1, 1, 1, N], [0, 0, 0, 1] dimensions are already special-cased in the compiler (and in some way in the hardware) to not use dimensions at all, so they would benefit from special case handling in IRON either way, I think.

@Victor-Jung also wrote some code to split up overflowing dimensions at some point IIRC.

[hunhoffe]

@andrej Hmm. I am realizing maybe I'm missing something in what you are proposing. Can you check my understanding against yours?

I was thinking that this should be handled in the compiler, not by the IRON python code. e.g., if [1, 1, 1, N], [0, 0, 0, 1] is already special cased in the compiler/hardware, I would think that the compiler should be able to catch whatever case the user (through IRON or mlir-aie generated through other means) throws at it and convert it into that special case is possible.

[andrej]

My perspective was that the "data layout transformation legalizations/optimizations" you describe are a distinct nice-to-have feature from "not using data layout transformations at all." From a user perspective:

• I can use data layout transformations.
  • Nice to have feature (what you're describing): the compiler automatically tries its best to legalize these transformations. Special case: The transformation ends up being linear/contiguous and the compiler detects this, opting to disable the use of the data layout transformation feature altogether.
• I can not use layout transformations (e.g. use a rt.fill(...) without a TAP).
  • If I do this with a 1D tensor type, IRON and the compiler generate a normal linear transfer, no issue.
  • If I do this with a (contiguous) ND tensor type (N > 1), IRON starts emitting code that uses the data layout transformation feature. As the user, I didn't express I wanted to use that feature (I may not be aware it exists). Even though we view this data as a ND-tensor, note that nothing changed about the transfer, it is still linear. Now, I might get errors that a data layout transformation I didn't express exceeds hardware limitations. This could be confusing and seems like a bug tome. Optimizing the transformations out again in a legalization pass seems like it could solve the problem, but it also seems like a bit of a roundabout solution when instead, if the user didn't express a layout transformation, we could simply not have emitted any code that uses the data layout transformation features in the first place (like we would have for the 1D case).

At the end of the day, it might be a pedantic distiction. To me, one just seems more like a bug and the other more like a feature. Fixing the bug seems like less work than implementing the (optional) feature.

Note also that some of the confusion may come from the fact of how weirdly we chose to encode "don't use data layout transformations" in MLIR. Rather than just say dimensions=None or something, it's the very specific [1, 1, 1, N], and [0, 0, 0, 1] sizes/strides that turn off the use of the data layout transformation feature altogether. That's why you can have N > 1024 in that case, but not any others.

[Victor-Jung]

I have a very simple code snippet to legalize taps in the mha operator in IRON. It's very simple and detects if the transfer is continuous, and if yes it tried to make it valid by using the [0, 0, 0, 1] stride tricks.