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

Add Support for ONNXCustomOp

Open chentong319 opened this issue 1 year ago • 4 comments

The current ONNXCallOp is used to call a function defined within ONNX model. When the function is not defined, namely is an external function, ONNXCustomOp is used. It turned out the ONNXCustomOp is quite useful to represent special operations when onnx flow graph is manipulated without defining a particular new Op. Easy for experiment and use of library. However, currently there is no support of shape inference and lowering for ONNXCustomOp.

Shape inference

Without the function body, we need extra info to perform shape inference. I think there are two possible approaches:

  1. Handle the common cases. Define an attribute to describe a common pattern, such as multi-dimension broadcast, same shape with a list of inputs, same or a particular type for element type. We can write the shape inference for ONNXCustomOp based on this specification. Hopefully, the common pattern can cover our use cases.
  2. Allow user to define particular shape inference function to associate it with the function name attribute in ONNXCustomOp. This approach is able to handle any case, but requires compiler modification.

Call convention

I found that there are two ways to pass the tensor to library call:

  1. Pass tensor as OMTensor data structure, which contains the data pointer, element type, and shape (dynamic) info.
  2. Pass only the data pointer for a tensor, and user will explicitly pass the shape info as extra parameters. User can use ONNXSizeOp or ONNXShapeOp to retrieve the shape info.

An attribute for the call convention will be added to the ONNXCustomOp and will be used when ONNXCustomOp is lowered to llvm IR.

Return method

Currently, an ONNX Op can be lowered to Krnl.Call, and then Krnl.Call is lowered to llvm function call. The major work from ONNX Op to Krnl.Call is to pass the output MemRef as a parameter, instead of returning a memory object allocated inside the library. This will complicate the memory management. However, is such requirement general enough? Do we need to specify which parameter will be used for output?

chentong319 avatar May 02 '23 23:05 chentong319