How do I pass bfloat16s to AIE application?

[giacomo-brunetta]

Hello everyone,
I am trying to write a Riallto application that performs floating point operations to a buffer of numbers. How do I modify this code so that the input and output are of the type bfloat16?

from npu.build.appbuilder import AppBuilder
import numpy as np
from npu.build.kernel import Kernel

def log_behaviour(invobj):
    invobj.out_buffer.array = np.array(shape = (1,1,1024), dtype = np.float32)

logv.behavioralfx = log_behaviour

class LogApplication(AppBuilder):

    def __init__(self, kernel:Kernel):
        self.kernel = kernel
        super().__init__()

    def callgraph(self, x_in, x_out) -> None:
        x_out[:] = self.kernel(x_in)

x_in = np.ones((1, 1, 1024), dtype=np.float32)
x_out = np.zeros((1, 1, 1024), dtype=np.float32)

app_builder = LogApplication(kernel=logv)
app_builder(x_in, x_out)

input_buffer = my_app.allocate(shape=(1,1,1024), dtype = np.float32)
output_buffer = my_app.allocate(shape=(1,1,1024), dtype = np.float32)

input_buffer[:] = np.ones((1,1,1024))*0.25

input_buffer.sync_to_npu()

my_app.call(input_buffer, output_buffer)

output_buffer.sync_from_npu()
print(output_buffer[0,0,0:16])

Numpy does not support bfloat16 and (in my understanding) tf and torch tensors are not a valid input type for AIE applications. Is there a work around to this problem?

[mariodruiz]

Hi @giacomo-brunetta,

Check this

https://github.com/AMDResearch/Riallto/blob/main/tests/test_alldatatypes.py#L10

[giacomo-brunetta]

Thanks, that was helpfull.
Writing from ml_dtypes import bfloat16 (as in the example) allows the use of bfloat16 as a datatype in Python and the data passing works correctly.

[mariodruiz]

if this solves the issue, please mark it as answered.