Use SSA values instead of integer BD IDs

[andrej]

We have some operations currently, such as npu_writebd, npu_push_queue and npu_dma_memcpy_nd that are a very thin wrapper around register writes that use integer values for buffer descriptor IDs.

I think there is an opportunity here to do something slightly smarter, and use SSA values with a new BD type instead. This SSA value would be returned from a configure operation, for example, or a dedicated allocation operation.

Why? It probably almost never matters what the actual integer value of the BD ID is. What does matter is that you reference the same BD across operations, e.g. you want to make sure that you push the same BD to the task queue using npu_push_queue as the one you just configured using npu_writebd.

Integers fail to make this easy; the compiler will happily emit code that pushes a BD ID to the task queue that is not configured.

With SSA values, we can make it impossible to push a BD to the queue that has not been configured.

Pros

• Mainly the above motivation.
• Can automatically allocate BD ids.
• Could also make it easier to piece code together from different places -- currently, if different code uses the same BD IDs, you need to rename them in one of the sources. With an allocator-based design, this wouldn't be an issue (as long as you don't run out of concurrently used BD IDs).

Cons

• Some things may be slightly harder to do, e.g. chaining BDs: If you want to chain a BD using next_bd currently, you can set this to a BD that you have not configured yet. When using SSA values, you would have to configure them in reverse order, so you have an SSA value available to reference the next BD. In the same context, we would have to think about how to handle cycles.

To do

• Also consider symbols? I don't know if symbols or SSA values are a more appropriate use in this case.
• Contemplate how to implement an allocator for BD IDs. It would be fairly simple, but needs to make sure no more than 16 BD IDs are live at the same time.
• Identify all ops that use integer BD ids
• Add a BD type
• Refactor all of those into using SSA values

[jgmelber]

Do you have a mock up of how this would look for a few use cases?

[jgmelber]

The BD chain syntax relies heavily on locks which we wouldn't always want to use when we are programming using the firmware because the "awareness" of an operation completing is localized to that lock and not visible to the firmware

[jgmelber]

It is also verbose to describe bd_chains with single BDs and push them to a channel's queue just to reuse the syntax when that is the most typical use case currently

[andrej]

As I understand it, you can leave off the locks if you don't want to use them.

This is probably one of those things that to a certain extent is a matter of taste. But I personally think the AIE syntax of configuring BDs is much easier to read than the long list of attributes to npu.writebd. In my opinion, we shoud emit a sequence of npu.writebd operations as the lowering of the above proposed "bd_chain".

Conceptually, configuring a BD somewhere on the array is very similar or identical to configuring BDs on the shim, except that on the shim we can re-configure dynamically. Since it's conceptually the same, I think the syntax should also be the same. This is why I am particular to this mock up.

Now we could talk about whether the current syntax in AIE is "good", but in any case I think it wouldn't hurt for it to at least be all the same everywhere. Once unified, we could revisit the current AIE syntax.

I do definitely agree that things would be quite a bit more verbose with this syntax. More typing, but also better readability, imo. This could also be alleviated with higher level abstractions.

Just my two cents! Happy to hear other arguments.

[stephenneuendorffer]

This is probably one of those things that to a certain extent is a matter of taste. But I personally think the AIE syntax of configuring BDs is much easier to read than the long list of attributes to npu.writebd. In my opinion, we shoud emit a sequence of npu.writebd operations as the lowering of the above proposed "bd_chain".

+1

[jgmelber]

This is probably one of those things that to a certain extent is a matter of taste. But I personally think the AIE syntax of configuring BDs is much easier to read than the long list of attributes to npu.writebd. In my opinion, we shoud emit a sequence of npu.writebd operations as the lowering of the above proposed "bd_chain".

+1

Yes, this sounds good! Looking forward to the working examples.

[andrej]

Here's the current proposal:

Defining a simple BD on the NPU shim

    aie.bd_chain @my_bd(%buf: memref<1xi32>) {
        aie.dma_bd(%buf : memref<1xi32>, 0, 1)
        aie.end
    }
    func.func @sequence(%arg0: memref<1xi32>, %arg1: memref<1xi32>) {
      %tkn1 = aiex.npu.start_bds @my_bd(%arg0) on (%tile00, MM2S, 0) {issue_token = 1}
      aiex.npu.wait(%tkn1)

Chain of BDs on the shim

    aie.bd_chain @my_chain(%bufA: memref<1xi32>, %bufB: memref<1xi32>) {
      ^bd0:
        aie.dma_bd(%bufA : memref<1xi32>, 0, 1, [<size = 1, stride = 2>, <size = 3, stride = 4>])
        aie.next_bd ^bd1
      ^bd1:
        aie.dma_bd(%bufB : memref<1xi32>, 0, 1, [<size = 1, stride = 2>])
        aie.end
    }
    func.func @sequence(%arg0: memref<1xi32>, %arg1: memref<1xi32>) {
      %tkn1 = aiex.npu.start_bds @my_chain(%arg0, %arg1) on (%tile00, MM2S, 0) {issue_token = 1}
      aiex.npu.wait(%tkn1)

1. The NPU lowering pass will generate buffer descriptor IDs automatically for those aie.dma_bd ops that do not explicitly specify one (such as the ones above). A global state will keep track of which BD IDs have been assigned on each NPU shim tile.
2. An aiex.npu.wait(%tkn1) will free all BD IDs of a previously issued BD chain for re-use. Alternatively, a aiex.npu.free_bd_ids(%tkn1) can be used to allow the allocator to re-use BD IDs that were assigned by that token. (This will not lower to any instruction but allow the compiler to reuse BD IDs.) (Useful when we can guarantee through other synchronization mechanisms that the BDs have completed.)
3. For now, only buffers and locks (things that are already SSA values in the current implementation) can be parametrized in the aie.bd_chain. In the future additional things may be come parametrizable (e.g. wraps/strides).
4. BD chains can be instantiated multiple times on different channels, or the same channel after synchronization.

Example Use of new BD Chains on aie.mem, aie.memtile_dma

    aie.bd_chain @ping_pong_pattern(%ping: memref<1xi32>, %pong: memref<1xi32>, %acq_lock: index, %rel_lock: index) {
      ^bb1: 
        aie.use_lock(%ping_acq_lock, AcquireGreaterEqual, 1)
        aie.dma_bd(%ping : memref<1xi32>, 0, 1)
        aie.use_lock(%ping_rel_lock, Release, 1)
        aie.next_bd ^bb2
      ^bb2:
        aie.use_lock(%pong_acq_lock, AcquireGreaterEqual, 1)
        aie.dma_bd(%pong : memref<1xi32>, 0, 1)
        aie.use_lock(%pong_rel_lock, Release, 1)
        aie.next_bd ^bb1
    }
    aie.mem(%tile02) {
      aie.dma_start_task @ping_pong_pattern (%fifo_output_buff_0, %fifo_output_buff_1, %fifo_output_cons_lock, %fifo_output_prod_lock) on (MM2S, 0)
    }

1. aie.expand_bds is a new operation. It would be a terminator of a block. In an initial pass, it will concretize the BD chain with the given arguments -- essentially copying and pasting the body of the BD chain while replacing the arguments with the concrete arguments provided. After expansion, lowering happens as before with the current passes (assign BD Ids etc.)
2. This would make writing certain patterns such as the ping-pong double buffering much less typing and clearer. We could potentially rewrite the objectfifo to generate patterns only once and then use aie.expand_bds with those patterns where needed to produce shorter and more readable IR.

Edit: Updated the last example to use a new syntax using aie.dma_start_task.

[andrej]

Next steps:

• Settle on final syntax and operation names after discussion with everyone; initially, let only buffer addresses and locks be "placeholders" through SSA values
• Implement parser and ops for the new syntax
• Implement lowering to NPU
• Implement lowering to HSA
• Expand new bd_chain operation (result of previous to-do) to be used outside of sequence function, in array aie.mem and aie.memtile_dma configuration, e.g. with a aie.dma_start_task operation. See last example above.
• (Possibly) extend what can be a placeholder in the bd_chain; e.g. make strides/wraps be an input parameter to the bd_chain op (Challenge: all inputs to dma_bd and use_lock are attributes, not SSA values. So could only allow constant SSA values.)