Fusing Initializers with Graph Transforms (microsoft#24726)

### Description

Added a graph transform for mixed precision graphs when FP16 compute is
unavailable. At session creation, this graph transform converts FP16
initializers (_which were changed to FP16 to FP32 cast nodes_) to FP32
initializers and fuses them with their next FP32 nodes.
 

- Behavior before this change:
"fp16 initializers -> cast_from_fp16_to_fp32 -> fp32 node/s"
 
- Behavior after this change:
"fp16 initializers converted to fp32 initializers then fused with fp32
node/s"

### Motivation and Context

This change aims to run the FP16 models without the repetitive casting
of FP16 initializers to FP32 initializers, by fusing FP32 initializers
with their next nodes, when FP16 compute is not available.

> For naming purposes, the newly added Graph Transforms in long form is
called "Fused Initializers Graph Transforms", and in short form is
called "FIGT".

### Working

Currently, the Fuse Initializers Graph Transform fuses cast nodes that
casts from FP16 to FP32, back to their
next/output nodes. Below is an explanation of how this transforms works.
It depends on ```InsertCastTransforms```
to produce the intermediate representation from which it fuses the
initializers (which are the cast node with
zero input, one initializer, and one output) back to the next/output
node. After fusion, the link/edge between such
a cast node to the next/output node will then be removed. Cast nodes
will be removed as well.

```
        "Input Graph"                       "Intermediate Representation"                 "FIGT Transforms"

          --------                   --------        --------        --------                 --------
         | X_Fp16 |                 | X_Fp16 |      | W_Fp16 |      | B_Fp16 |               | X_Fp16 |
          --------                   --------        --------        --------                 --------
             |                          |               |               |                        |
             |                          |               |               |                        |
             |                          V               V               V                        V
             |                       | Cast |        | Cast |        | Cast |                 | Cast |
             |                       | Fp16 |        | Fp16 |        | Fp16 |                 | Fp16 |
             |                       |  To  |        |  To  |        |  To  |                 |  To  |
             |                       | Fp32 |        | Fp32 |        | Fp32 |                 | Fp32 |
             |                          |               |               |                        |
             |                          |               |               |                        |
             V                          V               V               V                        V
 ----------------------------       -----------------------------------------       ----------------------------
|        Conv_Fp16           |     |                                         |     |         Conv_Fp32          |
|        --W_Fp16--          | ==> |                Conv_Fp32                | ==> |         --W_Fp32--         |
|        --B_Fp16--          |     |                                         |     |         --B_Fp32--         |
 ----------------------------       -----------------------------------------       ----------------------------
             |                                          |                                        |
             |                                          |                                        |
             |                                          V                                        V
             |                                       | Cast |                                 | Cast |
             |                                       | Fp32 |                                 | Fp32 |
             |                                       |  To  |                                 |  To  |
             |                                       | Fp16 |                                 | Fp16 |
             |                                          |                                        |
             |                                          |                                        |
             V                                          V                                        V
          --------                                   --------                                 --------
         | Y_Fp16 |                                 | Y_Fp16 |                               | Y_Fp16 |
          --------                                   --------                                 --------
```

The newly added Graph Transforms perform the following actions.

* Detect Cast node/s with single FP16 initializer converting to FP32.
* Convert all such FP16 initializer/s to FP32 initializer/s.
* Fuse newly created FP32 initializer/s to relative FP32 node/s.
* Remove FP16 to FP32 Cast node/s.

This is run in a loop as follows. It excludes Level 1 and Partitioning
optimizations.

```
 Level 2 --> Level 3 --> InsertCastTransforms --> FIGT
   ^                                        	   |
   |                 "LOOP"                	   |
   |                                       	   |
   -------------------------------------------------
```

### Adding FIGT as a Level-4 Graph Transform.

This will have the following benefits.

1. Ability to turn off (any/all) the Level 4 Optimizations. We can use
the `disable optimizers` functionality to turn off one of such
optimizations during testing, or use the `-o` switch to turn off all
Level 4 optimizations while executing a model using the command line or
Python scripts (or any other scripts).

2. Ability to rerun Level 2 and Level 3 optimizations remains intact
after Level 4 Optimizations are applied. Adding Level 4 takes care that
FIGT (or any similar optimizations) will always run after
InsertCastNodes.

3. It keeps the current graph manipulations untouched and gives us more
flexibility to add future optimizations like adding `Int8 to Int32`
upconvert or `FP8 to FP16` upconvert under Level 4. Level 4 can, as of
now, work as a placeholder for any other such upcoming Graph
optimizations.

```
 Level 2 --> Level 3 --> InsertCastTransforms --> Level 4
   ^                                        	    |
   |                 "LOOP"                	    |
   |                                       	    |
   --------------------------------------------------
```
> Added a placeholder for Level-4 for graph transforms utils under
orttraining. This helps resolve any exceptions that may be encountered
during training sessions.

#### Re-running Level 2+ optimizations after Level 4 / FIGT

The idea behind re-running Level2+ graph transforms is that, after the
fusion of initializers with their respective nodes, the nodes are now in
a format that might be supported by other graph transforms that were
previously skipped. Hence, some of the transformations previously unable
to be applied are now valid and can be applied to create a more optimal
graph for execution.

### Added a new session option
"kOrtSessionOptionsGraphOptimizationsLoopLevel" to handle the graph
optimization loop.

* When set to 2 or above it will loop until no more optimizations are
applied at any level starting Level 2 and above.

```
  Level 2 --> Level 3 --> InsertCastTransforms --> Level 4
    ^                                                 |
    |                 "Loop"                          |
    |                                                 |
    ---------------------------------------------------
```

* When set to 1 (default) it will loop until no more optimizations are
applied at Level 4 only.

```
  Level 2 --> Level 3 --> InsertCastTransforms --> Level 4
    ^                                                 |
    |        "Loop only depending on Level 4"         |
    |                                                 |
    ---------------------------------------------------
```

* When set to 0 it disables the loop.

```
  Level 2 --> Level 3 --> InsertCastTransforms --> Level 4
    ^                                                 |
    |                 "No Loop"                       |
    |                                                 |
    X                xxxxxxxxxxx                      X
```

### Documentation

We have not added any details related to Level 4 in the [Graph
Optimizations in ONNX
Runtime](https://onnxruntime.ai/docs/performance/model-optimizations/graph-optimizations.html)
documentation.

### OLD PR

This PR is created following a thorough discussion on the [OLD
PR](microsoft#24175).

Signed-off-by: Sunny Shukla <[email protected]>

Author: sunnyshu-intel

csharp/src/Microsoft.ML.OnnxRuntime/SessionOptions.shared.cs

@@ -17,6 +17,7 @@ public enum GraphOptimizationLevel
         ORT_DISABLE_ALL = 0,
         ORT_ENABLE_BASIC = 1,
         ORT_ENABLE_EXTENDED = 2,
+        ORT_ENABLE_LAYOUT = 3,
         ORT_ENABLE_ALL = 99
     }
 

include/onnxruntime/core/optimizer/graph_transformer_level.h

@@ -12,8 +12,9 @@ enum class TransformerLevel : int {
   Level1,       // basic optimizations
   Level2,       // extended optimizations
   Level3,       // layout optimizations
+  Level4,       // unsupported datatypes optimizations
   // The max level should always be same as the last level.
-  MaxLevel = Level3
+  MaxLevel = Level4
 };
 
 }  // namespace onnxruntime

include/onnxruntime/core/session/onnxruntime_c_api.h

@@ -355,6 +355,7 @@ typedef enum GraphOptimizationLevel {
   ORT_DISABLE_ALL = 0,
   ORT_ENABLE_BASIC = 1,
   ORT_ENABLE_EXTENDED = 2,
+  ORT_ENABLE_LAYOUT = 3,
   ORT_ENABLE_ALL = 99
 } GraphOptimizationLevel;
 

include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h

@@ -107,6 +107,37 @@ static const char* const kOrtSessionOptionsMemoryOptimizerProbeConfig = "optimiz
 // Default is an empty string which means no optimizers are disabled.
 static const char* const kOrtSessionOptionsDisableSpecifiedOptimizers = "optimization.disable_specified_optimizers";
 
+// It controls whether to run graph optimizations in loop or not.
+//
+// "0": disable. Graph Optimization Loop is disabled.
+// ```
+// Level 2 --> Level 3 --> InsertCastTransforms --> Level 4
+//   ^                                                 |
+//   |                 "No Loop"                       |
+//   |                                                 |
+//   X                xxxxxxxxxxx                      X
+// ```
+// "1": enable. Graph Optimization Loop is enabled, such that, if optimizations at Level 4 are applied then
+// the loop will check for any other valid optimization that can happen.
+// ```
+// Level 2 --> Level 3 --> InsertCastTransforms --> Level 4
+//   ^                                                 |
+//   |        "Loop only depending on Level 4"         |
+//   |                                                 |
+//   ---------------------------------------------------
+// ```
+// "2": enable. Graph Optimization Loop is enabled, such that, if optimizations at Level 2 or above are applied then
+// The loop will check for any other valid optimization that can happen.
+// ```
+// Level 2 --> Level 3 --> InsertCastTransforms --> Level 4
+//   ^                                                 |
+//   |                    "Loop"                       |
+//   |                                                 |
+//   ---------------------------------------------------
+// ```
+// Default value is set to "1".
+static const char* const kOrtSessionOptionsGraphOptimizationsLoopLevel = "session.graph_optimizations_loop_level";
+
 // Enable or disable using device allocator for allocating initialized tensor memory. "1": enable; "0": disable. The default is "0".
 // Using device allocators means the memory allocation is made using malloc/new.
 static const char* const kOrtSessionOptionsUseDeviceAllocatorForInitializers = "session.use_device_allocator_for_initializers";

java/src/main/java/ai/onnxruntime/OrtSession.java

@@ -652,6 +652,8 @@ public enum OptLevel {
        * graph.
        */
       EXTENDED_OPT(2),
+      /** Applies all the layout optimizations like NCHW and NCHWC to the ONNX graph. */
+      LAYOUT_OPT(3),
       /** Applies all available optimizations to the ONNX graph. */
       ALL_OPT(99);
 

java/src/main/native/OrtJniUtil.c

@@ -47,6 +47,8 @@ GraphOptimizationLevel convertOptimizationLevel(jint level) {
             return ORT_ENABLE_BASIC;
         case 2:
             return ORT_ENABLE_EXTENDED;
+        case 3:
+            return ORT_ENABLE_LAYOUT;
         case 99:
             return ORT_ENABLE_ALL;
         default:

js/common/lib/inference-session.ts

@@ -81,7 +81,7 @@ export declare namespace InferenceSession {
      *
      * This setting is available only in ONNXRuntime (Node.js binding and react-native) or WebAssembly backend
      */
-    graphOptimizationLevel?: 'disabled' | 'basic' | 'extended' | 'all';
+    graphOptimizationLevel?: 'disabled' | 'basic' | 'extended' | 'layout' | 'all';
 
     /**
      * Whether enable CPU memory arena.

js/node/src/session_options_helper.cc

@@ -31,6 +31,7 @@ const std::unordered_map<std::string, GraphOptimizationLevel> GRAPH_OPT_LEVEL_NA
     {"disabled", ORT_DISABLE_ALL},
     {"basic", ORT_ENABLE_BASIC},
     {"extended", ORT_ENABLE_EXTENDED},
+    {"layout", ORT_ENABLE_LAYOUT},
     {"all", ORT_ENABLE_ALL}};
 
 const std::unordered_map<std::string, ExecutionMode> EXECUTION_MODE_NAME_TO_ID_MAP = {{"sequential", ORT_SEQUENTIAL},

js/react_native/android/src/main/java/ai/onnxruntime/reactnative/OnnxruntimeModule.java

@@ -326,6 +326,7 @@ public WritableMap run(String key, ReadableMap input, ReadableArray output, Read
               {"disabled", SessionOptions.OptLevel.NO_OPT},
               {"basic", SessionOptions.OptLevel.BASIC_OPT},
               {"extended", SessionOptions.OptLevel.EXTENDED_OPT},
+              {"layout", SessionOptions.OptLevel.LAYOUT_OPT},
               {"all", SessionOptions.OptLevel.ALL_OPT},
           })
           .collect(Collectors.toMap(p -> (String)p[0], p -> (SessionOptions.OptLevel)p[1]));

js/react_native/ios/OnnxruntimeModule.mm

@@ -301,6 +301,7 @@ - (NSDictionary*)run:(NSString*)url
   @"disabled" : @(ORT_DISABLE_ALL),
   @"basic" : @(ORT_ENABLE_BASIC),
   @"extended" : @(ORT_ENABLE_EXTENDED),
+  @"layout" : @(ORT_ENABLE_LAYOUT),
   @"all" : @(ORT_ENABLE_ALL)
 };