GGML Operation Requirements

[ypapadop-amd]

Here's the list of requirements for supporting GGML operations:

1. A tensor is described by its data type, shape (4 dimensions), and strides in bytes on those dimensions. Since the stride is in bytes, it's derived from the data type. The stride is in bytes because they support sub-byte quantized types. There are no limitations in the shape or the strides. One can have a tensor that is 5005 elements (common in the test cases).
2. A lot of models are FP16. We convert automatically from FP16 to BF16.
3. Each operation has up to 10 input tensors and 1 output tensor. 4. The output may be an alias for any of the inputs. An input can be a view, i.e., an alias over another tensor with modified layout. E.g., in the transpose operation Tvw = T', Tvw is a tensor that aliases the data of T with modified strides to mimic iterating over a transposed matrix; Tvw may be passed to another operation, but not all operations support views as inputs.5. Some operations support additional invariants, e.g., a scale factor and a bias in the scale operation. These invariants are set dynamically during the operation instantiation, but they remain static after.
4. The data types of the inputs and/or the output are not necessarily the same.
5. For some operators, broadcast is supported in one or more of the inputs. E.g., doing C = A + B where A is a vector Mx1 and B is a matrix MxN replicates A on-the-fly to do a matrix addition without need for extra storage.

Additional requirements for integration:

1. Compilation artifacts are stored in directories that are dictated by the GGML library.
2. PDIs / insts files follow a specific naming scheme that is generated in the GGML library.
3. Compilation artifacts are kept in directories that match the name of the PDI file they produced.

[hunhoffe]

Thank you for writing this out! I won't make a concrete response to these issues until we get closer to being ready to implement them but this issue will ensure they remain on the radar!

[hunhoffe]

Some of these requirements may be closer on the roadmap, and some of them further. TL;DR -- 3 seems like something that could be immediately addressed, part of 4a seems like it could be immediately addressed, 4b seems like it should be prioritized as well. 1-3 integration requirements are also within the scope of work already planned for the immediate future. The others may need to be hacked/implemented as a top-level compatibility layer as it requires some further deep thought and time to discern the correct way to integrate them into the IRON framework.

My detailed thoughts:

• 1) Being a close-to-metal representation means that (for the most part, so far) the sizes/strides that are exposed in IRON almost directly reflect the sizes/strides supported by hardware, rather than reflecting the whole solution space of sizes/strides. It makes sense to me that GGML expects a different model, but I'm not sure whether the correct response is to change the model in IRON/mlir-aie or make a compatibility level as a top-level layer. My initial impression is that a proof-of-concept at the top level may be a good way to start. Some steps toward utilities for working with more generic sizes and strides are found here: https://github.com/Xilinx/mlir-aie/tree/main/python/helpers/taplib (with docs here: https://xilinx.github.io/mlir-aie/python/html/namespacetaplib.html). One day some additional compatibility layers could find there home there or within the compiler. For more immediate needs, if we were to design a minimal compatibility for the top level for an initial sample application, it might be nice to narrow the set of sizes/strides for now as a middle ground. @ypapadop-amd do you think the problem could be narrowed to simplify it for immediate needs?

• 2) I think BF16 is already fairly well supported in the existing operators here in IRON and the designs in mlir-aie. Is there a feature that is currently missing to support this, or is this mentioned just for completeness?

• 3) I think flexible numbers of arguments fit into our existing goals of having a generic base layer, and I think this should be prioritized. I believe there is already some work in this area towards combining multiple I/O buffers into one and automatically extracting them via offset in the runtime sequence. The first version of this will probably only allow concatenated data, and not interpolated data (e.g., sizes/strides). @andrej is this related to what you've already worked on?

• 4a) I think, as I mentioned it three, our goal is to support flexible arguments which includes arguments that are both input and output. However, the sizes/strides/views as inputs/outputs to operators is a more complex task that I believe will take some real thought. Let's revisit this after addressing some of the immediate TODOs; this type of capability is quite important to other tasks, such as composition/spatial pipelining, etc.

• 4b) Flexible use of datatypes across arguments should be supported as a priority. I don't think this will be terribly complex to implement, as we already do this with mlir-aie designs, including GEMM. @andrej Do you see any technical barriers here?

• 5) I think this needs to be left until after we get a solid base layer for operators, and I see this as a special case of the infering-sizes-strides-for-datamovement that would be required to address all of 4a, so my answer is the same as support for views/strided data as inputs/outputs.

I think all three of the other integration requirements are something we should minimally support in the basic definition of an operator.

@ypapadop-amd if there is anything you can do to narrow the solution space of sizes/strides that are most important to GGML that would be helpful to characterize the type of problems an initial compatibility layer may need. An initial first attempt at narrowing this space may be to only focus on the specific transformations/views required in a proof-of-concept application. I would appreciate your help with extracting a tractable initial implementation plan for these features!

[ypapadop-amd]

Thanks for your detailed reply.

• 1) Being a close-to-metal representation means that (for the most part, so far) the sizes/strides that are exposed in IRON almost directly reflect the sizes/strides supported by hardware, rather than reflecting the whole solution space of sizes/strides. It makes sense to me that GGML expects a different model, but I'm not sure whether the correct response is to change the model in IRON/mlir-aie or make a compatibility level as a top-level layer. My initial impression is that a proof-of-concept at the top level may be a good way to start. Some steps toward utilities for working with more generic sizes and strides are found here: https://github.com/Xilinx/mlir-aie/tree/main/python/helpers/taplib (with docs here: https://xilinx.github.io/mlir-aie/python/html/namespacetaplib.html). One day some additional compatibility layers could find there home there or within the compiler. For more immediate needs, if we were to design a minimal compatibility for the top level for an initial sample application, it might be nice to narrow the set of sizes/strides for now as a middle ground. @ypapadop-amd do you think the problem could be narrowed to simplify it for immediate needs?

It is difficult for non-experts to express complex access patterns mainly due to lack of examples. I wouldn't expect IRON to offer all the solutions, but a few non-trivial kernels with TAP would be most welcome, even if they are not the top performing.

I would expect any compatibility layer / glue layer to be at the GGML library, not in IRON, unless it's useful for other libraries. The conversation is mostly about how difficult it is to express computation with complex data movement. I think that we are leaving a lot of performance and abilities of the AIEs on the table just because of how obscure it is to describe such kernels.

• 2) I think BF16 is already fairly well supported in the existing operators here in IRON and the designs in mlir-aie. Is there a feature that is currently missing to support this, or is this mentioned just for completeness?

Most of simple GGML kernels uses FP32 (esp. for testing). There's a lot of FP16 as well, but not that much BF16, except specific models. I'll assume for performance reasons, IRON seems to be mostly BF16.

This creates a conundrum when it comes to working with anything else than BF16. I would expect that the full set of natively supported datatypes can be used with kernels (via templated kernels, datatype-configurable designs, etc.) to test operator correctness and model performance against CPUs/GPUs. Once this has been established, we can optimize for BF16 via further refining the model, etc. But it's very difficult to debug when working with different data types.

We would also need a kernel that converts FP16 to BF16 (this is currently done on the CPU). I'm looking into it based on the FP32 to BF16 conversion kernel, but if someone has already one, we could use it.

@ypapadop-amd if there is anything you can do to narrow the solution space of sizes/strides that are most important to GGML that would be helpful to characterize the type of problems an initial compatibility layer may need. An initial first attempt at narrowing this space may be to only focus on the specific transformations/views required in a proof-of-concept application. I would appreciate your help with extracting a tractable initial implementation plan for these features!

I can compile a list of a few kernels that can be used to run MNIST and some test kernels.

[andrej]

3. I think flexible numbers of arguments fit into our existing goals of having a generic base layer, and I think this should be prioritized. I believe there is already some work in this area towards combining multiple I/O buffers into one and automatically extracting them via offset in the runtime sequence. The first version of this will probably only allow concatenated data, and not interpolated data (e.g., sizes/strides). @andrej is this related to what you've already worked on?

Right now, the number of arguments is tied to XRT buffer objects and we're sometimes hitting limitations there -- apart from that, we already have that flexibility. The work I've started on subslicing arguments would further help with this indeed. Yes, striding in a larger buffer doesn't have an abstraction currently -- you have to manually implement that within the runtime sequence.

4b) Flexible use of datatypes across arguments should be supported as a priority. I don't think this will be terribly complex to implement, as we already do this with mlir-aie designs, including GEMM. @andrej Do you see any technical barriers here?

No immediate barriers come to mind. The dequant operator in the repository already uses a different data type for input/output.

@ypapadop-amd:

I would expect any compatibility layer / glue layer to be at the GGML library, not in IRON, unless it's useful for other libraries. The conversation is mostly about how difficult it is to express computation with complex data movement. I think that we are leaving a lot of performance and abilities of the AIEs on the table just because of how obscure it is to describe such kernels.

I think we touched on this in the llama meeting earlier with Erika. Let me share what I said there, I'd be curious to hear your thoughts.

I believe we first need to build a "layer zero" abstraction for operators, which abstracts only (1) the runtime management and (2) compilation of operators. We're relatively close to that.

Then, there are the tensor-related (complex data movement) abstractions: how do we express strides/wraps, transposes, data formats, etc. Getting a good abstraction for this should enable interoperability for tensor-related operators and, hopefully, also make writing the operators themselves easier. But I think we don't want to include them in the lowest layer -- layer zero should be generic, to e.g. also support string manipulation (bad example) operations who can still benefit from the runtime management and the compilation artifact logic, but don't need the tensor logic.

For your existing GGML integration, I believe, you have both: Logic for runtime management, compilation artifact generation, but also logic for tensor management/complex data movement. Right now, I don't think we're at the point where we can easily write operators that support arbitrary tensors with strides, transposes, etc. It'll remain a lot of work to (1) find the proper abstractions for this and (2) make the operator implementations use these abstractions.

I think the immediate path forward is to keep the necessary tensor wrangling logic inside GGML and operator implementations, but use the runtime/compilation artifact management from IRON so we can exercise and improve those parts. Concretely: The GGML layer can use IRON operators, but needs to pass the proper flags for data layout to each operator (e.g., read the metadata of the input tensor and deduce from that that it needs to pass b_col_maj=True to the GEMM operator as a parameter).

Then, as we learn how GGML and other users wrangle the tensors we can start standardizing those operator parameters into a unified "tensor" abstraction (e.g., GEMM's b_col_maj=True makes way for a more generic TensorDesc.get_layout() associated with its input argument that it can query).

Note that even once we have a proper tensor abstraction, I don't think a lot of operators will just "fall out" of that abstraction naturally, since the hardware wraps/strides, etc., are fairly restrictive and you have to be creative with your data movement patterns and hand-implement this. So, all work being done now to implement these operators is not lost -- they will only be "standardized" to accept a default tensor input format later.

So to summarize my thoughts: Step one is to make the "layer zero" abstraction (runtime/buffer management, compilation artifact generation) work for the GGML use case and use it in GGML. Step two is to implement the data wrangling needed to support arbitrary transposed inputs, strides, etc., as GGML requires them, by adding the necessary parametrizations/data movement patterns to the existing AIE operators. Step three is to learn from the GGML use case and other use cases as to which repeating patterns we see across these operators and unify them into a more general tensor abstraction that allows us to standardize the operator's parameters.

Most of simple GGML kernels uses FP32 (esp. for testing). There's a lot of FP16 as well, but not that much BF16, except specific models. I'll assume for performance reasons, IRON seems to be mostly BF16.

Indeed, bf16 is the only native vectorized non-block floating point data type in the AIE as far as I'm aware. Anything else is emulated.

Thoughts?

[ypapadop-amd]

Then, there are the tensor-related (complex data movement) abstractions: how do we express strides/wraps, transposes, data formats, etc. Getting a good abstraction for this should enable interoperability for tensor-related operators and, hopefully, also make writing the operators themselves easier. But I think we don't want to include them in the lowest layer -- layer zero should be generic, to e.g. also support string manipulation (bad example) operations who can still benefit from the runtime management and the compilation artifact logic, but don't need the tensor logic.

Agreed, level 0 to me sounds like it's a datatype agnostic, so tying it to a typical tensor representation may be counterproductive. However, I can see some value in maintaining a representation which can describe multiple data blocks within a buffer.

For your existing GGML integration, I believe, you have both: Logic for runtime management, compilation artifact generation, but also logic for tensor management/complex data movement. Right now, I don't think we're at the point where we can easily write operators that support arbitrary tensors with strides, transposes, etc. It'll remain a lot of work to (1) find the proper abstractions for this and (2) make the operator implementations use these abstractions.

We have those. There's also a possibility of adding extra logic to detect unexpected layouts and offer a sequence of kernels that should do the right thing.

However, to this date, I haven't been able to write a vector-add with different layouts in the input and output tensors; instead, temporary storage is allocated and the non-default layout tensors are copied using the CPU. I'd like to move away from this if possible.

Step one is to make the "layer zero" abstraction (runtime/buffer management, compilation artifact generation) work for the GGML use case and use it in GGML. Step two is to implement the data wrangling needed to support arbitrary transposed inputs, strides, etc., as GGML requires them, by adding the necessary parametrizations/data movement patterns to the existing AIE operators. Step three is to learn from the GGML use case and other use cases as to which repeating patterns we see across these operators and unify them into a more general tensor abstraction that allows us to standardize the operator's parameters.

I don't see step one as crucial: runtime/buffer management is tied to the ML framework and since every ML framework is opinionated about this it may be difficult to generalize it. For GGML it also doesn't matter, as we only need IRON for the kernels and the development environment. Once the kernels are compiled, IRON may not even be present, e.g., one big box w/ IRON but w/out an NPU compiles the kernels via an IRON environment, and multiple containers run in smaller boxes w/out IRON but w/ NPUs run inference.

On the other hand, utility kernels that transform data and data layouts are important until we figure out composition of data transformations w/ computation kernels.

[ypapadop-amd]

The current target are the following operations (MNIST-FC, all FP32):

1. X*Y: [784, 500], [784, 500] -> [500, 500] (GGML matmul is defined as C' = A x B')
2. x+y: [500, 500], [500, 1] -> [500, 500] (this does a broadcast to expand the second matrix)
3. RELU: [500] -> [500]
4. X*Y: [500, 10], [500, 500] -> [10, 500]
5. x+y: [10, 500], [10, 1] -> [10, 500] (broadcast here too)
6. argmax: [10, 500] -> [500, 1]
7. cross entropy loss: [10, 500], [10, 500] -> [1, 1]
8. cross entropy loss: [500, 1], [500, 1] -> [1, 1]

[hunhoffe]

• X*Y: [784, 500], [784, 500] -> [500, 500] (GGML matmul is defined as C' = A x B')

Have a C' = A x B' GEMM, but not for f32, and dimensions will need some padding. There is no native f32 operation on the AIE. Slow is ok for now though. Memcopy example has an example for padding, but this is not integrated for the GEMM yet. And instead of padding, you are re-using portions of data as padding (so not adding zeros).

Do we bucket + divisibility by 4 for padding or just make a point solution? There are some operations at the 4-byte level, which may make it hard.

Path forward for now: Point solution for this particular solution, but have the padding be handled by the operation internally (so no separate host operation to reorder). Use scalar operations for handling f32. We would like to handle more generic dimensions eventually, but we will start with the simplest thing we can get running.

• x+y: [500, 500], [500, 1] -> [500, 500] (this does a broadcast to expand the second matrix)

We have elementwise add, does all minus broadcast. Should be able to use one shim tile. Can we expand it?

Path forward for now: Yiannis is starting on this. Memory on L1 cores to store 1 500 element row of Y (broadcasted), stream X in one at a time, and then reuse that local copy of Y for each of the n X rows each core will process. Maybe look at the matrix vector TODO(erika)-IRON matrix-vector mul example.

• RELU: [500] -> [500]

We have this one. Double check dimensions.

• X*Y: [500, 10], [500, 500] -> [10, 500]

Path forward for now: Same path as for other GEMM.

• x+y: [10, 500], [10, 1] -> [10, 500] (broadcast here too)

Path forward for now: Same path as for other x+y matrix vector.

• argmax: [10, 500] -> [500, 1]

Elementwise operation? Max per row. Should hopefully be straightforward.

Path forward for now: Start with one of the IRON elementwise & modify it. Some elementwise take just one input, use one of those as a base. TODO(erika)-IRON example.

• cross entropy loss: [10, 500], [10, 500] -> [1, 1]
• cross entropy loss: [500, 1], [500, 1] -> [1, 1]

Softmax + reduction.
Can implement as one operator, or could chain two operators together. Reduction is not powers of 2, so dimensions are the harder part. MHA has some things via runtime parameter -- stream in power of 2 data, but use RTP to mask out data that should not be considered.

Path forward for now: This will likely be somewhat more challenging (along with the GEMM). Maybe try masking? Can be compile-time for point solution. TODO(erika)-find link to masking.

TODO: next steps may be aie2p vectorized f32 emulation for the kernels, but first step is just scalar.