onnx-mlir icon indicating copy to clipboard operation
onnx-mlir copied to clipboard

Include `onnx_node_name` in MLIR attributes for zhigh dialect ops

Open cjvolzka opened this issue 2 years ago • 1 comments

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>

image


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"}>>

cjvolzka avatar Jul 21 '22 16:07 cjvolzka

I believe this should be possible, there is also line info that can be embedded in each op.

AlexandreEichenberger avatar Jul 21 '22 21:07 AlexandreEichenberger