Using ::acquire_equal and ::release intrinsics for runtime-locking from inside kernel code instead of MLIR

[lstasytis]

We have been trying to construct an example where instead of calling acquire(ObjectFifoPort.Consume, 1) and release(ObjectFifoPort.Consume, 1) inside the python code which turns into MLIR, we use the ::acquire_equal and ::release instrinsics found inside <aie_api/aie.hpp>. The purpose of this being to enable data-driven locking behavior from inside the kernel code (we would want to have free-running kernels that we only call once for the entire execution of an application).

However we have ran into two issues in attempting this:

a.) the two intrinsics functions simply break our kernels, where a basic pass-through instead of copying the input array to the output array using a compute tile, doesn't move anything from simply calling either or both of these functions. We have tried using various lock ids and various values for what to set the locks to but nothing seemed to work. Is there any example we could use for how to use these intrinsics & confirmation that they should indeed work on the Ryzen NPUs and locking behavior is not limited to what can be generated for MLIR.

b.) Moving the locking to the kernel-side means also not using the acquire and release functions in python, however these seem necessary to turn an object fifo into a memref which can be passed to a kernel function call. From looking at the source code, we observed that a ObjectFifoSubviewType type object is created during lock acquisition using ObjectFifoSubviewAccessOp. Is there a way to just directly pass a memref to the kernel without calling the lock acquisition function?

Cheers,
Lukas

[AndraBisca]

Hello @lstasytis and thank you for your questions.

I don't currently have an answer for your first question, but regarding the use of locks with Object FIFOs, there are two possibilities to achieve what you propose:

The first is to set the disableSynchronization Object FIFO attribute to true, which will disable the creation of locks: this means that the acquire and release operations will not produce any MLIR involving locks, only to point to the correct aie.buffers; however, without synchronization, you may run into other issues with your data as it may be moved prematurely, so do watch out for that. This attribute is intended for debugging and is not exposed at the top-level of IRON, only at the placed variant. You can also generate the first level of MLIR and set the attribute to the Object FIFOs there.

The other option is to first produce the MLIR with Object FIFOs, then call the Object FIFO lowering pass to lower one step further using the aie-opt --aie-objectFifo-stateful-transform <name of MLIR file> command. From there you can delete any locks you do not need manually and use that MLIR code in your design.

Please let me know if either of these can work for you.

[AndraBisca]

Could you please also send the code for the kernel your wrote that uses lock intrinsics and what your target architecture is?

[lstasytis]

Hi @AndraBisca , thanks for the quick response!

Regarding the lock disabling for objectfifos question, disableSynchronization might work, but raises one question - could I set this attribute to disable the lock generation for one acquire and then turn it back on for another one?

I ask because I assume npu_dma_memcpy_nd performs a lock acquire & release from the side of the shim-tile and disableSynchronization might disable this as well, right? Or would an alternative then be to manually perform the locking from the shim tile?

The plan is to let the shim tile write to the memory tiles, start up the compute tile kernels and then be completely hands-off - let the kernels orchestrate all the work between themselves with the locks internally and only check for a result lock to be released to then perform the writes back to DDR from the memory tile. If this can be done closer to the metal using something like the kernel described in this test we would happily use that, but the single unit test is not much to go off of.

Regarding the locks inside the kernels question, after further tests we did manage to get the below kernel.cc code to run and produce a proper pass-through. This is specifically a 0,0 -> 0,1 -> 0,2 -> 1,2 -> 1,1 -> 1,0 topology of routing data from 0,0 shim to 1,0 shim using the 0,2 and 1,2 compute tiles. However this only partially works. We are targeting aie2p in this case. (npu2 - we are using one of the MiniPCs). For vitis essentials, we are using 1.3 ryzen software.

#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <type_traits>
#include <aie_api/aie.hpp>

extern "C" {

void tile_0_2(int32_t *a, int32_t *c, int32_t N) {

  //::acquire_equal(18,0);
  for (int i = 0; i < N; i++) {
    c[i] = a[i];
  }
  //::release(18,1);
}

void tile_1_2(int32_t *a, int32_t *c, int32_t N) {
  // block until lock 50 has been incremented
  ::acquire_greater_equal(50,1);
    for (int i = 0; i < N; i++) {
    c[i] = a[i];
  }
  ::release(50,1);

}

} // extern "C"

For this example, we adapted the section-3 of the programming guide example, introduced a second set of shim, memory and compute tiles and then removed the lock acquire and release calls such as

mlir-aie/programming_guide/section-3/aie2.py

Line 51 in 7e75546

elem_in = of_in.acquire(1)

(to still get the memref, we moved them to before the for-loop and just did an acquire->release two-liner).

While this works, performing the same with the tile_0_2 no longer manages to proceed (kernel_1 fails to acquire the lock) and leads to no passthrough. (It doesn't hang, actually, just produces all 0s, interestingly). If I don't remove the mlir-side of the locks, uncommenting those two lines does produce a result and only fails to passthrough the final for loop iteration.

The biggest source of confusion is how to even determine what the lock ids should be between the tiles. I got the 18 and 50 in this case from looking at the MLIR generated using the objectfifo acquire and release functions in the python code. I assume there is some system that would allow to deterministically map adjacent tile lock IDs, but I don't see it being documented anywhere. We would really use a few examples of using these locks in the kernels.

P.S while on the issue: what exactly happens to the MLIR-defined locks when the MLIR code and the kernel code are linked and then executed on the device? Are the locking instructions appended to the start and end of the kernel programs? Or are the locks set by a master controller/host? Curious what the trade-off here is outside of just a different style of programming.

[jgmelber]

For vitis essentials, we are using 1.3 ryzen software

Any reason you're using Chess rather than Peano?

[lstasytis]

For vitis essentials, we are using 1.3 ryzen software

Any reason you're using Chess rather than Peano?

We're experimenting with both, but our understanding has been that Peano's support for instrinsics is a little flaky, while xchesscc should be much more reliable. Coming from an FPGA background, we want complete control of the hardware, so something as close to the programming model as what is provided for the Versals (at least until Peano has more support for general purpose computing other than ML workloads (read: data-dependent execution flow))

[AndraBisca]

Regarding the lock disabling for objectfifos question, disableSynchronization might work, but raises one question - could I set this attribute to disable the lock generation for one acquire and then turn it back on for another one?

I ask because I assume npu_dma_memcpy_nd performs a lock acquire & release from the side of the shim-tile and disableSynchronization might disable this as well, right? Or would an alternative then be to manually perform the locking from the shim tile?

The disableSynchronization feature directs to compiler to not generate any locks for the synchronization of that Object FIFO. That means that all acquire and release calls will be affected, on both producer and consumers(s). However, the Object FIFO will not generate any configuration for the shim tile, unless the register_external_buffers() function was used. This is because the shim tile interacts with memory locations outside of the tile-array and the Object FIFO only allocates memory within it. The Object FIFO attributes and their lowering will not affect the npu_dma_memcpy_nd operations.

[jackl-xilinx]

@lstasytis Lock ids are positionally encoded depending on which tile is referencing which lock. Every tile has 16 local locks with a tile's local memory being directionally to it's east. The lock id's to these 16 locks associated with its local memory is lock id 48-63 (hence lock id 50 might be chose to represent lock 2). To access the locks for neighboring local memories, the indexing is as follows:

South:  0-15   0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15
West:  16-31  16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31
North: 32-47  32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47
East:  48-63  48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63

As you can see, the tile to the right who wishes to refer to the same lock id would refer to the west encoding so lock 2 for him would be lock id 18. If you reverse the direction, then the references would change. Technically, the locks used doesn't have to match the local memory being accessed. All that's important is that both neighboring tiles reference the same set of locks. If you look at the build/*mlir.prj/input_physical.mlir, you'll see the mlir encoding for lock designation like:
%in_cons_prod_lock_0 = aie.lock(%tile_0_2, 0) {init = 2 : i32, sym_name = "in_cons_prod_lock_0"}

This sets a lock for tile_0_2_0 to be lock 2 with the symbolic name in_cons_prod_lock_0. That would normally be lock id 50. As long objfifos and lowering steps manages these but if you put them in your kernel code and leverage objfifos lowerings, then you would have to confirm the lock id used in the lowering and then make sure you refer to the right id value for the tile where your kernel is running.

[jgmelber]

Maybe #2898 is helpful here 😄

[lstasytis]

Maybe #2898 is helpful here 😄

Extremely, thank you, @jgmelber ! We'll be playing around with the locks in the coming weeks and let you know if anything odd pops up.

[jgmelber]

I am still playing around with it from the objectFIFO API side but the lock functionality is there