onnx-mlir
onnx-mlir copied to clipboard
Include `onnx_node_name` in MLIR attributes for zhigh dialect ops
This is a request for enhancement to the MLIR output. When --EmitMLIR
is used with NNPA acceleration, it would be useful to have some sort of reference in the MLIR output to indicate which onnx.* a particular zhigh.* replaces.
As an example, currently when compiling bertsquad-12 for NNPA, some MatMul Operations went to CPU and some were lowered to NNPA.
The CPU (ie onnx.* Ops) include onnx_node_name
in the attributes dictionary. That can be copy and pasted into a netron.app view of the model to make it easy find the Op in a graphical view.
%317 = "onnx.MatMul" (%316, %299) {onnx_node_name = "bert/encoder/layer_0/attention/self/MatMul_1"} : (tensor<?x?x256x256xf32>, tensor<?x?x?x?xf32>) -> tensor<?x?x256x?xf32>
Later when another MatMul Op is lowered to NNPA in MLIR, There's no `onnx_node_name`. So it is not as easy to identify exactly which MatMul a particular zhigh operation is replacing.
%323 = "zhigh.MatMul"(%320, %321, %322): (tensor<?x?xf32, #zhigh.encoding<{dataLayout = "2D"}>>, tensor<768x768xf32, #zhigh.encoding<{dataLayout = "2D"}>>, tensor<768xf32, #zhigh.encoding<{dataLayout = "1D"}>>) -> tensor<?x768xf32, #zhigh.encoding<{dataLayout = "2D"}>>
I believe this should be possible, there is also line info that can be embedded in each op.