Support for multiple individual tasks in the same file

[ipapal]

Hello,

I would like to ask whether/how Aries supports having multiple tasks within the same file, that I can handle them individually.

I’ve successfully followed the provided examples and tutorials to run multiple GEMM operations in a single file. However, I’m interested in managing each operation independently, particularly with regard to custom scheduling parameters. Specifically, I’d like to apply different parallelization strategies to each GEMM operation. Is this currently supported by Aries?

As far as I have understood, with my current flow all GEMM operations map to the same AIEs resources, while what I would like to do, is have different AIEs allocated per GEMM task. I’ve attempted to achieve this by defining multiple functions using the @task_tile decorator, but I consistently encounter compilation errors when running the Python script.

This also extends to supporting different AIE kernels in the same file. With a simple example in which I try to have two diverse GEMM AIE kernels, the compilation fails.

Best regards

[adhowe3]

From what I have heard, ARIES does not support multiple @task_kernel() or @task_tile() declarations currently (beta.v0). But they plan on implementing this in a future release.

I also had tried doing this to add bias after a matrix multiply, but it does not seem possible right now. I ended up implementing the bias addition inside the same kernel. Something like this:

@task_kernel()
def kernel_gemm(TileA: float32[TI, TK], 
                TileB: float32[TK, TJ], 
                TileC: float32[TI, TJ], 
                tileBias: float32[TJ]):
    for i0 in range(0, TI):
        for j0 in range(0, TJ):
            TileC[i0, j0] = float32(0)
            for k0 in range(0, TK):
                TileC[i0, j0] += TileA[i0, k0] * TileB[k0, j0]
            TileC[i0, j0] += tileBias[j0]

I don't know if that is helpful or not. Hopefully, one of the team members can weight in on this issue.