Implementing Relu in ARIES - unsupported operators max(), comparators (>, <, <=, >=)

[adhowe3]

I am working on implementing a full Res-Net model using ARIES. So far I have a convolutional layer, bias addition, and batch normalization layers. It has been tricky, but I have been able to make it work all in a single kernel.

I cannot implement a ReLU (Rectified Linear Unit) function using ARIES, because there is not support for a max() function or >. This is a simple function that returns 0 for negative inputs and the original value for positive inputs. ReLU is usually implemented with one of the following:

return max(x, 0)

return x * (x > 0)

It looks to me like the only current operations that generate MLIR from ARIES python are *, /, +, - as seen in aries_ir_builder.py

It looks like in the Aries/lib emitter files there is support for a max() function and many other operations not yet supported on the frontend.

I was able to modify the generated AIE code to produce the results of a max() function. It was rather simple, but doing it by hand is not ideal for large projects:

v51 = fpmac(v51, v52, 0, 0x76543210, v50, 0, 0x76543210);  // previous operation that I want to pass through ReLU
aie::vector<float,8> relu = max(v51, v17);  //  ReLU operation

[JinmingZhuang]

Hi,

Thank you for reaching out.

In the playground of ARIES, there is an example for IR Builder Example. By running this file, it will print out the ARIES initial IR and users could explore adding new operations starting from it. Currently, the generated MLIR for single kernel is shown below but is not as expected:

func.func @kernel_relu(%TileA: memref<32xf32, 2>, %TileC: memref<32xf32, 2>) attributes {adf.kernel} {
affine.for %i0 = 0 to 32 {
    %v0 = affine.load %TileA[%i0] : memref<32xf32, 2>
    affine.store %v0, %TileC[%i0] : memref<32xf32, 2>
    %c0 = arith.constant 0.000000 : f32
    affine.store %c0, %TileC[%i0] : memref<32xf32, 2>
  }
  return
}

This is because in current IR builder, the "if" and "compare" nodes in Python AST haven't been supported yet.

To add other operations, users could modify the IR builder for the single kernel here Kernel builder. To view the Python AST of the single kernel, you could uncomment this line Print Python AST.

Please refer to this page for details about Python AST Description. Then proper parsers should be implemented (e.g. visit_If, visit_Compare in Parser) to support Python-to-MLIR translation for "if" and "compare" nodes.