nod-ai / iree-amd-aie

IREE plugin repository for the AMD AIE accelerator
Apache License 2.0
45 stars 23 forks source link

Create global and local allocations #485

Open yzhang93 opened 4 days ago

yzhang93 commented 4 days ago

The goal is to create local allocations which could be a reference to the global allocation for each tile to access the memory. The prototype is as below

We leave the memref.alloc op as a global variable: %alloc = memref.alloc() : memref<1x1x8x4x8x4xi32, 2 : i32>

And create a new pointer pointing to the global allocation as %ptr = amdaie.addressof %alloc : memref<1x1x8x4x8x4xi32, 2 : i32> -> !amdaie.ptr

To load the memory: %local_alloc = amdaie.load %ptr : !amdaie.ptr -> memref<1x1x8x4x8x4xi32, 2 : i32>

And then replace all uses of %alloc with %local_alloc.

@MaheshRavishankar Could you take a look and see if I understand this correctly? And do I miss anything? @jtuyls It would be helpful if you could add some context of how we are going to use the local allocations on different tiles.

MaheshRavishankar commented 4 days ago

So we dont technically need a "global variable". We can leave the alloc as it is today, just move it to the top of the function.

And create a new pointer pointing to the global allocation as %ptr = amdaie.addressof %alloc : memref<1x1x8x4x8x4xi32, 2 : i32> -> !amdaie.ptr

you dont need a new type. You can just use

%memref_ref = amdaie.reference_to %alloc : memref<..., 2: i32>

%memref_ref is the same type as %alloc, but the reference is not hoistable across the scf.forall boundary. This will make the amdaie.load isntruction unnecessary. You just have to replace all %alloc with %memref_ref.