Memory Tile for Data Parallel Applications with custom kernels

[GiuseppeSorrentino99]

Hello everyone,
I want to use a custom kernel taking RTP params in 4 different compute tiles in the same column. I managed to use the MT tile for a single CT and followed what's done in notebook 4.7, which presents the MT usage for data-parallel applications.
However, the first one is assigned only, as if I am giving as input the same instance kernel.
I also tried to declare four different kernels with different names...but other errors occur.

I am probably missing something explained in the tutorials or straightforward, but I would deeply appreciate support with this.
I thank you all in advance, and I write the used code below in case one wants to try this.

This is the kernel:

%%kernel

#define vector_size 32 // Size of the vector unit, used for reading from input stream
#define read_type uint8 // type of the input stream
#define func_type int // type used for doing the computations
#define write_type float // type of the output stream

inline aie::vector<func_type,vector_size> convert_to_func_type(aie::vector<read_type,vector_size> vec) {
    aie::accum<acc32, vector_size> acc;
    acc.from_vector(vec, 0);
    return acc.to_vector<func_type>();
}

void partialmse(uint8_t* img1, uint8_t* img2, float* out, uint32_t size) {
    aie::vector<uint8_t,32> vect1;
    aie::vector<uint8_t,32> vect2;
    aie::vector<func_type,vector_size> diff;
    float32 partial_sum = 0;
    for(int i = 0; i< size/vector_size; i++){
        vect1 = aie::load_v<vector_size>(img1);
        img1+=vector_size;
        vect2 = aie::load_v<vector_size>(img2);
        img2+=vector_size;
        diff = aie::sub(convert_to_func_type(vect1), convert_to_func_type(vect2));
        auto square = aie::mul(diff, diff).to_vector<int>();
        partial_sum += aie::reduce_add(square);
    }
    write_type res = (write_type) partial_sum / size;
    out[0] = res;
}

And this is its behavioural FX

def partialmse_behaviour(invobj):
    invobj.out.array = np.ndarray(shape = (1,1,1), dtype = np.float32)

partialmse.behavioralfx = partialmse_behaviour

Finally, this is the entire application:

from npu.build.appbuilder import AppBuilder
import numpy as np
from npu.build.kernel import Kernel
from npu.build.mtkernel import MTSplit, MTConcat


class MT3MSEApplication(AppBuilder):

    def __init__(self, kernelA:Kernel,kernelB:Kernel, kernelC:Kernel, kernelD:Kernel):
        #if len({id(kernelA), id(kernelB), id(kernelC), id(kernelD)}) != 4:
        #    raise ValueError("Need 4 different kernel instances")

       self.kernel0 = kernelA 
       self.kernel1 = kernelB
       self.kernel2 = kernelC
       self.kernel3 = kernelD
       self.mtbsplit1 = MTSplit(4)
       self.mtbsplit2 = MTSplit(4)
       self.mtbconcat = MTConcat()
       super().__init__()

    def callgraph(self, x_in1:np.ndarray, x_in2:np.ndarray, x_out:np.ndarray):
    
        row_num = x_in1.shape[0]
        col_num = x_in1.shape[1]
        num_tiles = 4
        partial_outputs = [None] * num_tiles # I need partial outputs from split, to concatenate them
        size = col_num/num_tiles
        size = int(size)
        print(type(x_in1))
    
        for row in range(row_num):
            #split the input in 4 
                inputs1 = self.mtbsplit1(x_in1[row])
                inputs2 = self.mtbsplit2(x_in2[row])
                partial_outputs[0] = self.kernel0(inputs1[0], inputs2[0], size)
                partial_outputs[1] = self.kernel1(inputs1[1], inputs2[1], size)
                partial_outputs[2] = self.kernel2(inputs1[2], inputs2[2], size)
                partial_outputs[3] = self.kernel3(inputs1[3], inputs2[3], size)
                mtbuffer_out = self.mtbconcat(partial_outputs)
                x_out[row] = mtbuffer_out
    
appBuilder = MT3MSEApplication(kernelA=partialmse,kernelB=partialmse,kernelC=partialmse,kernelD=partialmse)

In this case, the four kernels should be four instances of the same AIE kernel (but they may be different ones), but this Application uses a single compute tile since it is seen as assigning the same instance to all the kernels.
Furthermore, I cannot "create" the specific kernels in the init since I should give an available RTP param in the callgraph.

Thanks for your support,

[GiuseppeSorrentino99]

As an additional information:

If I declare 4 different partialmse kernels ( 4 AIE kernel with different names), it can see 4 different instances, but I got the following error

 '<' not supported between instances of 'NoneType' and 'NoneType' 

but none of my compontens in the callgraph is None, and the other error hints are in mlir code

[mariodruiz]

@GiuseppeSorrentino99 have you tried with a list comprehension?

[partialmse for _ in range(4)], you do not need necessarily need to pass the kernel as argument to the constructor. Maybe try this.

Other thing that you can try is creating the kernel using the Kernel class.

[GiuseppeSorrentino99]

Thanks for your clarification:
I also did some tests in this direction:

I tried with two compute tiles only, but I got the same issue
'<' not supported between instances of 'NoneType' and 'NoneType'

Furthermore, If I have understood correctly, we have 6 inDMA and 6 outDMA. Is this correct?
If this is the case, and I am correctly counting the input and output ports:

**About the DMAin**
2 DMAin are required for the MTSplit inputs + 4 DMAin are the inputs for the concat. = 6

**About the DMAout**
 8 DMAout are from the MTSplits + 1 DMAout from the MTConcat = 9

Therefore, I also tried to remove one MTSplit. The CT Kernel can take the same input twice, coming from the MTSplit1
Doing this, I have

**About the DMAin**
1 DMAin isrequired for the MTSplit input + 4 DMAin are the inputs for the MTConcat. = 5

**About the DMAout**
4 DMAout are from the MTSplit + 1 DMAout from the MTConcat = 5

This configuration has the same issue, with both 2 and 4 CTs.

Thansk again for your support, I deeply appreciate it. Let me know if you want me to do further tests

[mariodruiz]

Furthermore, If I have understood correctly, we have 6 inDMA and 6 outDMA. Is this correct?

Yes, I suggest you check the notebook 3_2. All the details of the architecture are there.

I need to check the issue in more detail

[GiuseppeSorrentino99]

Thanks mario! I will surely deep notebook 3.2 for further details about the architecture.

In the meanwhile, let me know if you have any hints about this specific issue, and I will work on it as soon as possible.

[mariodruiz]

How do you build the app? can you also provide that part?

[GiuseppeSorrentino99]

Sure,

vol = 1
dim = 64
dimR = dim*4
dimC = dim*4
img1 = np.zeros(shape=(dimR,dimC,vol),dtype=np.uint8)
img2 = np.zeros(shape=(dimR,dimC,vol),dtype=np.uint8)
out = np.zeros(shape=(dimR, 2,1,1),dtype=np.float32) 

appBuilder(img1,img2,out)
appBuilder.build(img1,img2,out)

Notice that the output shape is (dimR, NumCTiles,1,1), hence in this code I am targeting 2 num_tiles ( the test you suggested me some hours ago).

[mariodruiz]

@GiuseppeSorrentino99,

Can we try to simplify the problem? Could you write a kernel that takes an array of uint8 and generates the float output? With this we can try to generate a similar callgraph and see id we narrow down the issue.

[GiuseppeSorrentino99]

Ciao Mario.
I think I follow your suggest, but still the same issue. Here is the code...maybe you can see some error?

This is the kernelgenerator.py ( the same you have on the repo)

from npu.build.kernel import Kernel

class KernelObjCall():
    """ Helper function that generates new kernel objects."""
    def __new__(cls, cpp, bfx, *args):
        kobj = Kernel(cpp, bfx)
        print(type(kobj))
        return kobj(*args) if len(args) > 0 else kobj

Then, I have a python class, similar to the one you have on the repo

from kernelgenerator import KernelObjCall
from pathlib import Path
import numpy as np

class SimplerPartialMSE():

    def __new__(cls, *args):
        cpp = str(Path(__file__).parent / "cpp" / "simplerpartialmse.cpp")
        val = KernelObjCall(cpp, cls.behavioralfx, *args)
        return val
        
    def behavioralfx(self):
        self.out.array = np.ndarray(shape = (1,1,1), dtype = np.float32)

While this is the kernel C for the AIE ( same to before, but it reads twice from the same input, as the images for mse are interleaved)

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


#define vector_size 32 // Size of the vector unit, used for reading from input stream
#define read_type uint8 // type of the input stream
#define func_type int // type used for doing the computations
#define write_type float // type of the output stream

inline aie::vector<func_type,vector_size> convert_to_func_type(aie::vector<read_type,vector_size> vec) {
    aie::accum<acc32, vector_size> acc;
    acc.from_vector(vec, 0);
    return acc.to_vector<func_type>();
}

// this version takes as input the two images interleaved
void simplerpartialmse_aie(uint8_t* img1 = nullptr, uint8_t* img1 = nullptr, float* out = nullptr, uint32_t size = 1) {
    aie::vector<uint8_t,32> vect1;
    aie::vector<uint8_t,32> vect2;
    aie::vector<func_type,vector_size> diff;
    float32 partial_sum = 0;
    for(int i = 0; i< size/vector_size; i++){
        vect1 = aie::load_v<vector_size>(img1);
        img1+=vector_size;
        vect2 = aie::load_v<vector_size>(img1);
        img1+=vector_size;
        diff = aie::sub(convert_to_func_type(vect1), convert_to_func_type(vect2));
        auto square = aie::mul(diff, diff).to_vector<int>();
        partial_sum += aie::reduce_add(square);
    }
    write_type res = (write_type) partial_sum / size;
    out[0] = res;
}


extern "C" {

void simplerpartialmse(uint8_t* img1, float* out , uint32_t size) {
    simplerpartialmse_aie(img1, img2, out, size);
}

} 

Following, instead, I put the code for the Application class

import npu
import numpy as np
from npu.build.appbuilder import AppBuilder
from npu.build.mtkernel import MTSplit, MTConcat
from PythonWrappers import SimplerPartialMSE

class MT3MSEApplication(AppBuilder):

    def __init__(self, numTiles):
        super().__init__()
        if(numTiles < 2 or numTiles > 4):
            raise ValueError("Pick Between 2 or 4 ")

        self.num_tiles = numTiles
        self.kernels = [SimplerPartialMSE() for _ in range(self.num_tiles)]
        self.mtbsplit1 = MTSplit(4)
        self.mtbconcat = MTConcat()

    def callgraph(self, x_in1:np.ndarray, x_out:np.ndarray):
        print("Start Callgraph")
        row_num = x_in1.shape[0]
        col_num = x_in1.shape[1]
        partial_outputs = [None] * self.num_tiles # I need partial outputs from split, to concatenate them
        size = col_num/(2*self.num_tiles)
        size = int(size)
        print(row_num)
        print(size)
        print(col_num)
        for row in range(row_num):
                inputs1 = self.mtbsplit1(x_in1[row])
                for i in range(self.num_tiles):
                    partial_outputs[i] = self.kernels[i](inputs1[i], size)
                    #print("Boom") PASSED
                print(type(partial_outputs))
                print(type(partial_outputs[0]))
                print(partial_outputs[0].shape)
                mtbuffer_out = self.mtbconcat(partial_outputs)
                x_out[row] = mtbuffer_out
        print("Exit Callgraph") # This print is correctly executed
        
n_tiles = 2
appBuilder = MT3MSEApplication(numTiles = n_tiles)

And this is when I build the program

vol = 1
dim = 64
dimR = dim*4
dimC = dim*4
img1 = np.zeros(shape=(dimR,dimC,vol),dtype=np.uint8)
img2 = np.zeros(shape=(dimR,dimC,vol),dtype=np.uint8)
out = np.zeros(shape=(dimR, 2,1,1),dtype=np.float32) 
appBuilder(img1,out)
appBuilder.build(img1,out) # HERE IT FAILS

I will conduct further tests on this side.

For now I have tried to:

1. Introduce errors in the kernel code -> it gives a different error. Hence I believe this is not the case.
2. I have tried with a very basic kernel, taking a single input of uint8 and a single output:

extern "C" {

void simplerpartialmse(uint8_t* img1, uint8_t* out , uint32_t size) {
    out[0] = 42;
}

} 

This has the same error.
There must be something wrongly translated in mlir...
cause the python application class completes each instruction.

Let me know if there are further tests I can try, and thanks again for the effort in helping me! I deeply appreciate it

[mariodruiz]

Posted a solution in #66