Correctly representing MIGraphX values in lowering

[krzysz00]

Note: doing this depends on getting #1140 , the general perf key ticket, done, because it'll be really ugly if we don't do that first

The problem

Currently, the MLIR migraphx dialect maps MIGraphX's shapes (which are {sizes,...}, {strides, ...}, type to upstream's tensor, which throws away stride information.

This hasn't been much of a problem historically, as MIGraphX would insert the appropriate transpose operations to reorder the shape to be in increasing-stride order ... on each input.

However, MIGraphX IR expects operation, like convolution, to preserve the strides. So, if MIGraphX wants an NHWC convolution (which is represented in IR with shape = {N, C, H, W}, stride = {CHW, 1, CW, C}), it expects the output layout to match the input layout.

However, because our representation of MIGraphX IR doesn't include stride information, we unconditionally produce an NCHW output.

This problem isn't limited to convolutions, though, and would affect any MIGraphX kernel request that has an output with a shape that isn't in "stride order".

This is, as far as I'm concerned, a historical mistake we made where we didn't design the MIGraphX dialect to properly represent the MIGraphX IR.

Proposed solution

While a short-term fix (inserting more migraphx.transpose ops onto outputs) has been identified, this entire decision of using tensor<> for MIGraphX's shape is likely to have irritating long-term impacts.

Therefore, to substantially simplify our code and make these edge cases much easier to handle, I propose the following changes (after #1140)

#migraphx.shape<L1xL2...xLk, S1xS2...xSk, T>

All operations in the migraphx dialect will no longer operate on the tensor type, but will instead use the migraphx.shaped<Sizes, Strides, Type> type. This will be the type MIGraphX generates when it translates its IR into an MLIR module.

This will also allow us to ensure we're correctly preserving every detail of MIGraphX's broadcasting semantics.

Function arguments

The one awkward thing here is function arguments, since those'll need to go through a generally unmodified MLIR bufferization flow. There, I propose a solution that'll also have useful effects for the underlying codegen:

%t = migraphx.mlir.arg_view %argN : tensor<actual_underlying_size_of_object()xT> to #migraphx.shape<L, S, T>
// and
%ret = migraphx.mlir.result_view %V : #migraphx.shape<L, S, T> to tensor<actual_underlying_size_of_object()xT>

That is, the arguments to the function become the underlying float *, half *, ... that you'll be passing in, and then, at the beginning of the function, we put the shape information back, and, similarly, the "returned" float * and friends are represented as the underlying memory and then we write it in the expected pattern.

During lowering

During MIGraphX to Tosa conversion, we use MLIR's type conversion system to insert the appropriate conversions.
Specifically,

• mlir.arg_view and mlir.result_view become the relevant rock.transform[Embed{}] to produce a tensor whose shape is what Tosa expects (the logical shape) but which writes in the underlying strided form. (In a Rock-less lowering, this would be memref.reinterpret_cast once someone defined a way to sneak it through tensor)
• With all other MIGraphX ops, the extra shape information will help us handle implicit broadcasts on the way to Tosa.
• As far as I can tell, we'll never encounter shape changes outside of function boundaries or an operation in the MIGraphX IR that explicitly fiddles with the shape (but, for example, transpose preserves the stride order, I think) so we'll be able to convert everything to tensor<LxT> on the way to Tosa

@sjw36 @pfultz2 @kahmed10 @manupak

[sjw36]

I am generally opposed to any remaining MIGraphX MLIR dialect after conversion to TOSA.

1 simpler change that would not affect bufferization, etc. would be to add the #migraphx.shape attribute to the function parameter attributes. These could then be applied after bufferization to produce the proper form. This would avoid any new operations.

Please add an example of the MIGraphX IR input to the problem statement so we can explore all alternatives.

[krzysz00]

The problem here is that I want to represent MIGraphX IR correctly so that we can't forget about strides during the translation process.

Here's an example of what we should be able to deal with (taken from NHWC work with the implicit transpose operations we currently add removed)

mlir_main:pointwise0:y1.0 = @param:y1.0 -> float_type, {2, 16, 3, 3}, {144, 1, 48, 16}, target_id=0
mlir_main:pointwise0:y0.0 = @param:y0.0 -> float_type, {1, 16, 4, 4}, {256, 1, 64, 16}, target_id=0
mlir_main:pointwise0:@4 = convolution[padding={0, 0, 0, 0},stride={1, 1},dilation={1, 1},group=1,padding_mode=0](mlir_main:pointwise0:@y1,mlir_main:pointwise0:@y0) -> float_type, {1, 2, 2, 2}, {8, 1, 4, 2}, target_id=0
mlir_main:pointwise0:@5 = exp(mlir_main:pointwise0:@4) -> float_type, {1, 2, 2, 2}, {8, 1, 4, 2}, target_id=0
mlir_main:pointwise0:@6 = @return(mlir_main:pointwise0:@5), target_id=0

[krzysz00]

Annotating function parameters isn't sufficient, since the stride information has to propagate to results

[krzysz00]

And I'll note that, with the conversion to Tosa, we could, instead of leaving any migraphx operations in, go straight to rock.transform during the translation process

[sjw36]

Annotating function parameters isn't sufficient, since the stride information has to propagate to results

Results can also have attributes.

And I'll note that, with the conversion to Tosa, we could, instead of leaving any migraphx operations in, go straight to rock.transform during the translation process

This is certainly preferred.

[krzysz00]

Ok, then we'll translate the operation I'm thinking to call migraphx.mlir.parameter_to_shaped and migraphx.shaped_to_result to a rock.transform during MIGraphX-to-Tosa

[manupak]

Sorry, I think this was discussed when Im out -- hence missed this.

A design choice to consider :

I think it is fair to use 'tensor' dialect alongside TOSA. If so we should be able to do a more direct lowering to https://mlir.llvm.org/docs/Dialects/TensorOps/#tensorextract_slice-tensorextractsliceop that preserves the semantics you are after. It would be then supported by https://github.com/ROCmSoftwarePlatform/rocMLIR-internal/issues/870.

There is also https://mlir.llvm.org/docs/Dialects/TensorOps/#tensorinsert_slice-tensorinsertsliceop

And I'll note that, with the conversion to Tosa, we could, instead of leaving any migraphx operations in, go straight to rock.transform during the translation process

This is certainly preferred.

This is going to break the clone-based tests.

[krzysz00]

Yeah, having considered this further, in most cases, we'll be able to go to tosa.reshape and tosa.transpose - I described it in the tickets. (and we need to just not have e2e tests for the really weird cases)

[krzysz00]

Closing because we've implemented this