[onnx/ort] Convert transformer optimizations using custom-ops to functions for optimizer/rewriter usage

[shubhambhokare1]

There are a few commonly used transformer-based optimizations using custom-ops in onnxruntime some of which are listed.

Custom ops to be converted to functions for ONNX Script (will add more custom ops based on findings and priority) :

Tasks

Beta Give feedback

1. GroupQueryAttention
2. RotaryEmbedding
3. SimplifiedLayerNormalization
4. SkipSimplifiedLayerNormalization
5. MatMulNBits
6. SparseAttention

The goal is to create an infrastructure via which these custom-ops can be stored as functions for ease of usage with the optimizer and rewriter tools. If the functions are not desired in the graph, the graph can be simply inlined to get back to the original form.

Road-map for this task:

Tasks

Beta Give feedback

1. Find the correct function signature/definition to be defined for each custom op. A good starting point could be work already done in the function-rewriter.
2. Create a common method to convert these custom-ops to functions
3. Add ONNX Spec for these custom ops so that they could be added to ONNX standard

@gramalingam Feel free to edit/add any more information that I might have missed regarding this item.

[gramalingam]

LGTM, thanks! What we want to do is:

1. Define these as ONNX functions (and test/validate them)
2. Extend the ONNX standard by adding these as function-ops
3. Potentially use these in a rewriting-optimizer (that identifies matches to the function-body and extracts them out as either model-local functions or standard ONNX ops, if they have been added to the ONNX standard as in step 2).

For steps 1 and 2, the corresponding custom op definition in ORT is also a helpful starting point.