Packet header modification from within tile kernel code

[lstasytis]

Me and a set of colleagues are looking at utilizing the NPUs for more complex control flow algorithm acceleration, coming from an FPGA background. In the process, behavior which we need is the ability to control packet switching at runtime from within the compute tiles.

A minimal example would be a chain of two compute tiles, each sending a packet back and forth, incrementing the value stored within until one of tiles detects that the packet's payload has increment up to some threshold value, at which point the tile would update the packet's header id to point to the shim tile such that the tile is routed back to the host, finishing the program's execution.

For this example to work, we need an ability to manipulate packet headers from within the CPP/Python APIs. Is there something preventing support of this feature?

Alternatively, we would use the Vitis AIE kernel APIs, however it is not clear if it is possible to target the Ryzen NPU tiles using that compiler. From our understanding, the Ryzen AI Software 1.3 suite ships with 'Vitis AIE Essentials' which may provide these APIs to target the NPUs with xchesscc?

Any clarification on this would be extremely helpful!

[fifield]

A minimal example would be a chain of two compute tiles, each sending a packet back and forth, incrementing the value stored within until one of tiles detects that the packet's payload has increment up to some threshold value, at which point the tile would update the packet's header id to point to the shim tile such that the tile is routed back to the host, finishing the program's execution.

For this example to work, we need an ability to manipulate packet headers from within the CPP/Python APIs. Is there something preventing support of this feature?

This is possible to do. The connections between the tiles and between the tile and shim can be packet routed, but the DMAs do not have to be configured to attach packet headers on the outgoing data. Instead you can create the packet header as part of the data buffer and have the tile program manage the packet header. The packet header format is found here with this parity function.

Here is an example that ping pongs a buffer between tiles then sends it to a shim tile after a threshold: https://gist.github.com/fifield/c8438132bf7903fa2b91fac01a50bfd4

[lstasytis]

Here is an example that ping pongs a buffer between tiles then sends it to a shim tile after a threshold: https://gist.github.com/fifield/c8438132bf7903fa2b91fac01a50bfd4

Thanks, @fifield , this is already immensely helpful to know that the behavior is supported on the MLIR front.

The follow-up question would then be: how involved would it be to add functionality to either the cpp or python APIs to emit the header read/write behavior in MLIR? We would very much prefer to not have to manually write kernels in MLIR but take advantage of all the bindings you've developed so far. Even something hacky would be fine, since we really only need this packet id manipulation to support complex execution flow at runtime.

Lastly, is using the payload for delivering the header really the only way? 15 elements means we can't use 16-element SIMD operations anymore without padding which will introduce its own set of headaches.

[fifield]

The follow-up question would then be: how involved would it be to add functionality to either the cpp or python APIs to emit the header read/write behavior in MLIR? We would very much prefer to not have to manually write kernels in MLIR but take advantage of all the bindings you've developed so far. Even something hacky would be fine, since we really only need this packet id manipulation to support complex execution flow at runtime.

The example I provided is just setting bits in buffers, and would be the same in C++ or Python as in MLIR. This is fairly bare metal and won't compose well with abstractions like object fifo but it doesn't require any special support. What sort of API or abstraction did you have in mind?

[lstasytis]

The example I provided is just setting bits in buffers, and would be the same in C++ or Python as in MLIR. This is fairly bare metal and won't compose well with abstractions like object fifo but it doesn't require any special support. What sort of API or abstraction did you have in mind?

Given that the main purpose of this functionality is to simply introduce control mechanisms between tiles, I don't think it has to be anything overly sophisticated for starting out.

If I understand the behavior correctly, this seems like it could be abstracted as a pair (or single?) function along the lines of

value = packet_header_get(packet_ref, byte_indx)
packet_header_set(packet_ref, byte_indx, value)
where,
packet_ref is the received buffer memory reference (%current_val = memref.load %buf_0_2_in[%c0] : memref<16xi32>)
byte_indx is the byte position to modify (or read if we don't drop the header automatically?)
value is the byte value to write at position byte_indx to packet_ref.

My concerns would be:
a.) that the MLIR has 7 instructions for setting these bits in the header which, if it really translates to 7 instructions on the compute tile, could destroy performance. The setting of these headers should ideally overlap with the compute so that a packet can be loaded, the data operated on with 1-2 instructions (like SIMD additions) and immediately sent to the next tile with a packet header informing which direction the packet should go to.

b.) that this is all only for packets travelling along the NoC and not the fast lanes between adjacent tiles. My understanding has been that the NoC is primarily for sending these sorts of control signals between tiles and is 32-bit wide, while actual data should be going through the ping pong buffers of adjacent tiles using 512 bit-wide ports. But since your example had 16-element payloads, that seemed to be the same sort of direct neighbor sending? The way I see all this working in real applications is that the slow NoC is largely ignored and packets are only traveling the the ping pong buffers neighbor-to-neighbor, with the packet headers informing the tile whether to forward it south/north/east/west or keep it. The NoC may be used for additional control signals to manage the state of the compute tiles. Let me know if my understanding is correct here.

Nevertheless, some sort of barebones function to simply avoid having to manually edit MLIR code for modifying the packets would be immensely helpful, even if it isn't the most efficient thing in the world yet, since we mainly want to unblock our thesis students right now so that they can start implementing these more complex algorithms on the devices. Optimizing everything to really squeeze every last bit of performance out of the devices would come next.