Common API - Network Definition API

[music-dino]

Create an MGX API to mimic the TRT network definition APIs.
Once the APIs are created, attempt to combine it with the TRT ONNX parser so that the network definition and layer manipulation APIs can be used on a network created by parsing an ONNX file.

[music-dino]

Input

• ITensor * addInput (char const *name, DataType type, Dims const &dimensions)

Wraps an mgx::parameter

kACTIVATION

• IActivationLayer * addActivation (ITensor &input, ActivationType type)

Activations which require non-default alpha and/or beta params, must set them after the layer is created via the appropriate method
Type can be changed after the layer is created.
ActivationTypes that map to a single MGX operator:

• kRELU - "relu"
• kSIGMOID - "sigmoid"
• kTANH - "tanh"
• kLEAKY_RELU - "relu"
• kELU - "elu"

ActivationTypes that map to a combination of MGX operators:

• kSELU
• kSOFTSIGN
• kSOFTPLUS
• kCLIP
• kHARD_SIGMOID
• kTHRESHOLED_RELU

Unsure:

• kGELU_ERF
• kGELU_TANH

Not part of ONNX:

• kSCALED_TANH

kLRN

• ILRNLayer * addLRN (ITensor &input, int64_t window, float alpha, float beta, float k)

Maps to a single MGX operator - "lrn"
WindowSize, Alpha, Beta, and K can be changed after the layer is created

kSCALE

• IScaleLayer * addScale (ITensor &input, ScaleMode mode, Weights shift, Weights scale, Weights power)

Doesn't seem to have ONNX equivalent

kSOFTMAX

• ISoftMaxLayer * addSoftMax (ITensor &input)

Axis must be set after the layer is created
Maps to a single MGX operator - "softmax"

kCONCATENATION

• IConcatenationLayer * addConcatenation (ITensor *const *inputs, int32_t nbInputs)

Axis must be set after the layer is created
Maps to a single MGX operator - "concat"

kELEMENTWISE

• IElementWiseLayer * addElementWise (ITensor &input1, ITensor &input2, ElementWiseOperation op)

Operation can be changed after the layer is created
ElementWiseOperations mapping to mgx operators:

• kSUM - "add"
• kPROD - "mul"
• kMAX - "max"
• kMIN - "min"
• kSUB - "sub"
• kDIV - "div"
• kPOW - "pow"
• kAND - "logical_and"
• kOR - "logical_or"
• kXOR - "logical_xor"
• kEQUAL - "equal"
• kGREATER - "greater"
• kLESS - "less"

ElementWiseOperations without mapping to mgx operators:

• kFLOOR_DIV - based on description, could be done by a combination of "div" and "floor"

kUNARY

• IUnaryLayer * addUnary (ITensor &input, UnaryOperation operation)

Operation can be changed after the layer is created
UnaryOperations mapping to mgx operators:

• kEXP - "exp"
• kLOG - "log"
• kSQRT - "sqrt"
• kRECIP - "recip"
• kABS - "abs"
• kNEG - "neg"
• kSIN - "sin"
• kCOS - "cos"
• kTAN - "tan"
• kSINH - "sinh"
• kCOSH - "cosh"
• kASIN - "asin"
• kACOS - "acos"
• kATAN - "atan"
• kASINH - "asinh"
• kACOSH - "acosh"
• kATANH - "atanh"
• kCEIL - "ceil"
• kFLOOR - "floor"
• kERF - "erf"
• kNOT - "not"
• kSIGN - "sign"
• kROUND - "nearbyint"
• kISINF - "isinf"

kSHUFFLE

• IShuffleLayer * addShuffle (ITensor &input)

This layer shuffles data by applying in sequence: a transpose operation, a reshape operation and a second transpose operation. The dimension types of the output are those of the reshape dimension.

The layer has an optional second input. If present, it must be a 1D Int32 shape tensor, and the reshape dimensions are taken from it.

Maps to a combination of mgx operators (transpose and reshape).

kONEHOT

• IOneHotLayer * addOneHot (ITensor &indices, ITensor &values, ITensor &depth, int32_t axis)

Axis can be changed after the layer is created
Maps to a combination of mgx operators

kREDUCE

• IReduceLayer * addReduce (ITensor &input, ReduceOperation operation, uint32_t reduceAxes, bool keepDimensions)

Operation, Axes, and KeepDimensions can be changed after the layer is created
ReduceOperations mapping to mgx operators:

• kSUM - "reduce_sum"
• kPROD - "reduce_prod"
• kMAX - "reduce_max"
• kMIN - "reduce_min"
• kAVG - "reduce_mean"

kTOPK

• ITopKLayer * addTopK (ITensor &input, TopKOperation op, int32_t k, uint32_t reduceAxes)

Operation, K, and reduceAxes can be changed after the layer is created
TopKOperations kMAX and kMIN correspond to the two possible values of the largest attribute?
setInput(1) is used when K is to be provided as an input. MGX does not support this

kGATHER

• IGatherLayer * addGather (ITensor &data, ITensor &indices, int32_t axis)

GatherMode must be set after the layer is created, unless kDEFAULT is desired
GatherAxis and NbElementWiseDims can be changed after the layer is created
Depending on GatherMode, the layer corresponds to different ONNX ops:

• kDEFAULT - Gather(deprecated)
• kELEMENT - GatherElements
• kND - GatherND

For mode kND the mgx "gathernd" operator can be used directly
For mode kELEMENT the "gather" operator can be used in combination with other operators(NOTE: There seems to be a bug in parse_gather_elements.cpp:95, 0 is always passed for axis)

• IGatherLayer * addGatherV2 (ITensor &data, ITensor &indices, GatherMode mode)

See above

kRAGGED_SOFTMAX

• IRaggedSoftMaxLayer * addRaggedSoftMax (ITensor &input, ITensor &bounds)

No equivalent ONNX or MGX operator

KMATRIX_MULTIPLY

• IMatrixMultiplyLayer * addMatrixMultiply (ITensor &input0, MatrixOperation op0, ITensor &input1, MatrixOperation op1)

Op0 and Op1 can be changed after the layer is created
MatrixOperations kNONE and kTRANSPOSE can be supported, kVECTOR needs more investigation
Can be implemented as GEMM with alpha and beta equal to 1

kNON_ZERO

• INonZeroLayer * addNonZero (ITensor &input)

Implemented via "nonzero" mgx op if input is not a constant, a literal is inserted otherwise

kCONSTANT

• IConstantLayer * addConstant (Dims const &dimensions, Weights weights)

Wraps an mgx literal
Dimensions and weights can be changed after layer is created

kIDENTITY

• IIdentityLayer * addIdentity (ITensor &input)

Directly maps to mgx id operator

kCAST

• ICastLayer * addCast (ITensor &input, DataType toType)

toType can be changed after the layer is created
Maps to mgx "convert" operator

kPARAMETRIC_RELU

• IParametricReLULayer * addParametricReLU (ITensor &input, ITensor &slope)

Maps to mgx "prelu" operator

kCONVOLUTION

• IConvolutionLayer * addConvolutionNd (ITensor &input, int64_t nbOutputMaps, Dims const &kernelSize, Weights kernelWeights, Weights biasWeights)

Input, nbOutputMaps, kernelSize, kernelWeights, and biasWeights can be changed after the layer is created
Dilations, nbGroups, PaddingMode, PaddingNd, PostPadding, PrePadding, StrideNd all must be set after the layer is created if non-default values are to be used
If setPaddingMode and setPrePadding/setPostPadding are both used, PaddingMode takes precedence
The supported TRT padding modes are:

• kEXPLICIT_ROUND_DOWN
• kEXPLICIT_ROUND_UP - one of the two explicit modes should correspond to the NOTSET ONNX mode
• kSAME_UPPER - corresponds to ONNX SAME_UPPER
• kSAME_LOWER - corresponds to ONNX SAME_LOWER
• There is no VALID padding mode. It would seem that a valid convolution is performed if PaddingMode, PrePadding and PostPadding are left at their default values(kEXPLICIT_ROUND_DOWN, zeroes, and zeroes respectively), but this needs to be tested out.

PrePadding and PostPadding amount to the same things as the pads attribute in ONNX. The difference being that ONNX prohibits pads being used when auto_pad is used, while TRT gives precedence to auto_pad/PaddingMode if both are used.

kPOOLING

• IPoolingLayer * addPoolingNd (ITensor &input, PoolingType type, Dims const &windowSize)

Corresponds to ONNX MaxPool, AveragePool, GlobalMaxPool and GlobalAveragePool
For global pooling a window size equal to the input size should be used
PoolingType kMAX_AVERAGE_BLEND does not correspond to an ONNX operator
LpPooling is not supported
PoolingType, and window size can be changed after the layer is created
AverageCountExcludesPadding, PrePadding, PostPadding, PaddingMode, and Stride must be set after the layer is created if non-default values are to be used
For padding information see addConvolutionNd

kDECONVOLUTION

• IDeconvolutionLayer * addDeconvolutionNd (ITensor &input, int64_t nbOutputMaps, Dims kernelSize, Weights kernelWeights, Weights biasWeights)

Corresponds to ONNX ConvTranspose operator
Setting the output dimensions is not supported in TRT
For other detail see addConvolutionNd

kSCALE

• IScaleLayer * addScaleNd (ITensor &input, ScaleMode mode, Weights shift, Weights scale, Weights power, int32_t channelAxis)

TODO

kRESIZE

• IResizeLayer * addResize (ITensor &input)

Implemented as a combination of mgx operators
OutputDimensions, Scales, ResizeMode, CoordinateTransformation, SelectorForSinglePixel, NearestRounding, CubicCoeff, ExcludeOutside must be set after the layer is created if non-default values are to be used
InterpolationMode corresponds to ONNX mode, all are supported
ResizeCoordinateTransformation corresponds to ONXX coordinate_transformation_mode, only half_pixel, asymmetric and align_corner are supported
ResizeRoundMode corresponds to ONNX nearest_mode, all are supported
Antialising is not supported

kLOOP

• [ ] ILoop * addLoop ()

TODO

kCONDITION

• IIfConditional * addIfConditional ()

TODO

kSELECT

• ISelectLayer * addSelect (ITensor &condition, ITensor &thenInput, ITensor &elseInput)

Equivalent to the "where" mgx/ONNX op

kASSERTION

• IAssertionLayer * addAssertion (ITensor &condition, char const *message)

No equivalent ONNX or MGX operator

kFILL

• IFillLayer * addFill (Dims const &dimensions, FillOperation op, DataType outputType)

Covers the RandomNormal and RandomNormalLike ONNX operators for KillOperation::kRANDOM_NORMAL, and RandomUniform and RandomUniformLike for FillOperation::kRANDOM_UNIFORM.
Value of the seed attribute, if present, is ignored
In MGX a literal with values from the appropriate distribution is created
Operation and outputType can be changed after the layer is created
Alpha and beta values must be changed after the layer is created if non-default values are to be used
Alpha and beta have different meaning depending on FillOperation:

• For UNIFORM, alpha is the minimum value, beta the maximum
• For NORMAL, alpha is the mean, beta is the standard deviation

kPADDING

• IPaddingLayer * addPaddingNd (ITensor &input, Dims const &prePadding, Dims const &postPadding)

Pre and post padding can be changed after the layer is created
Only does zero padding

kDEQUANTIZE

• IDequantizeLayer * addDequantize (ITensor &input, ITensor &scale, DataType outputType)

TODO

kSCATTER

• IScatterLayer * addScatter (ITensor &data, ITensor &indices, ITensor &updates, ScatterMode mode)

Covers the ONNX ScatterElements and ScatterND operators for ScatterMode::kElement and ScatterMode::kND respectively
Reduction is not supported
Mode can be changed after the layer is created
Axis must be set after the layer is created if non default value is to be used

...

[music-dino]

The initial approach will be to wrap both layers and tensors around instructions.
Layers that are implemented by more than one mgx operator will hold references to the starting and ending instructions of the block of instructions that it's composed of, while the output tensor(s) will hold references to only the final instruction.
When a layer is modified after creation by use of a setter method, the existing mgx IR will be modified by use of replace_instruction.

This approach allow us to use existing infrastructure for both instruction modification, which needs to be reflected in downstream instructions if the output type changes, and for creating composite operators, by reusing existing parsing code.

A difficulty that presents itself for any approach and this one especially are ONNX operators that hold subgraphs.
An attempt will be made to validate if this approach can work for them by implementing the ILoopLayer.
A good idea for how this layer works can be found in the trt onnx parser source code for parsing an ONNX Loop: https://github.com/onnx/onnx-tensorrt/blob/7ecb49a435bd881b9ac4011450315192885e5cc3/onnxOpImporters.cpp#L2810

[music-dino]

The way INetworkDefinition tracks input and output tensors needs to be modified. MGX does not guarantee that output parameters will be inserted after input parameters, nor seemingly the relative ordering of the parameters.