[Rewriter] Adopt missing ONNX Light optimization patterns

docs/articles/index.md

@@ -4,4 +4,5 @@
 :maxdepth: 1
 
 2023-06-20-introducing-onnx-script/index
+onnx_light_patterns
 ```

onnxscript/rewriter/__init__.py

@@ -37,17 +37,26 @@
 )
 from onnxscript.rewriter.rules.common import (
     _basic_rules,
+    _boolean_optimizations,
     _broadcast_to_matmul,
+    _canonical_optimizations,
     _cast_constant_of_shape,
     _collapse_slices,
+    _collection_optimizations,
+    _expand_optimizations,
     _fuse_batchnorm,
     _fuse_pad_into_conv,
     _fuse_relus_clips,
+    _label_encoder,
+    _layout_optimizations,
     _materialize_reshape_shape,
+    _matrix_optimizations,
     _min_max_to_clip,
     _no_op,
     _redundant_scatter_nd,
+    _remove_expand_before_binary_op,
     _remove_optional_bias,
+    _shape_optimizations,
 )
 
 logger = logging.getLogger(__name__)
@@ -62,6 +71,15 @@
     *_min_max_to_clip.rules,
     *_fuse_relus_clips.rules,
     *_basic_rules.basic_optimization_rules(),
+    *_boolean_optimizations.rules,
+    *_canonical_optimizations.rules,
+    *_collection_optimizations.rules,
+    *_shape_optimizations.rules,
+    *_layout_optimizations.rules,
+    *_expand_optimizations.rules,
+    *_remove_expand_before_binary_op.expand_before_binary_op_rules,
+    *_matrix_optimizations.rules,
+    *_label_encoder.rules,
     *_redundant_scatter_nd.rules,
     *_fuse_pad_into_conv.rules,
     *_fuse_batchnorm.rules,

onnxscript/rewriter/_ir_utils.py

@@ -52,6 +52,8 @@ def visit(node: ir.Node, depth):
 
 
 def get_const_value(value: ir.Value) -> ir.TensorProtocol | None:
+    if value.is_graph_input():
+        return None
     node = value.producer()
     if node is not None:
         optimizer.basic_constant_propagation([node])
@@ -61,11 +63,7 @@ def get_const_value(value: ir.Value) -> ir.TensorProtocol | None:
 def get_numpy_value(val: ir.Value | None) -> np.ndarray | None:
     """Convenience wrapper to get (optional) numpy value from an optional IR Value.
 
-    This is intended for use in optimizations/rewriting. Note that this does not
-    yet handle the distinction between inputs with default values (values that are
-    both graph inputs and graph initializers), which should not be treated as a
-    constant, and true constant values. The caller should make the distinction, as
-    a value does not contain enough information to determine this. (TODO)
+    Graph inputs with initializer defaults remain overridable and are not constants.
     """
     if val is None:
         return None

onnxscript/rewriter/ort_fusions/cos_sin_cache.py

@@ -38,6 +38,7 @@ def __init__(
         cast: bool = False,
         reshape: bool = False,
         const_freqs: bool = False,
+        concat_after_trig: bool = False,
     ):
         # This pattern makes use of shared Cos/Sin values. So, we can't remove the
         # matched nodes as part of the rewrite-step. We apply a separate final
@@ -50,6 +51,7 @@ def __init__(
         self._reshape = reshape
         self._cast = cast
         self._const_freqs = const_freqs
+        self._concat_after_trig = concat_after_trig
 
     @property
     def max_pos_id(self) -> int | None:
@@ -123,10 +125,12 @@ def pattern(
             # B: batch size, S: sequence length, E: embedding dimension
             # position_ids: [B, S] or [S]
             # inv_freq: [1, E, 1]
-            position_ids_expanded = op.Unsqueeze(
-                position_ids, extra_dims
-            )  # [B, S] | [S] => [B, 1, S]
-            position_ids_expanded = op.Cast(position_ids_expanded, to=ir.DataType.FLOAT)
+            position_ids_expanded = pattern.OrValue(
+                [
+                    op.Cast(op.Unsqueeze(position_ids, extra_dims), to=ir.DataType.FLOAT),
+                    op.Unsqueeze(op.Cast(position_ids, to=ir.DataType.FLOAT), extra_dims),
+                ]
+            )
             # if self._reshape:
             #     position_ids_expanded = op.Expand(position_ids_expanded, _allow_other_inputs=True)
             #     position_ids_expanded = op.Reshape(position_ids_expanded, _allow_other_inputs=True)
@@ -141,13 +145,16 @@ def pattern(
         # if self._reshape:
         #     freqs = op.Reshape(freqs, freqs_3d_shape)  # redundant reshape
         freqs = op.Transpose(freqs, perm=[0, 2, 1])  # [B, S, E]
-        emb = op.Concat(freqs, freqs, axis=-1)
+        emb = freqs if self._concat_after_trig else op.Concat(freqs, freqs, axis=-1)
         cos = op.Cos(emb)
         if self._cast:
             cos = op.Cast(cos, to=dtype)
         sin = op.Sin(emb)
         if self._cast:
             sin = op.Cast(sin, to=dtype)
+        if self._concat_after_trig:
+            cos = op.Concat(cos, cos, axis=-1)
+            sin = op.Concat(sin, sin, axis=-1)
         cos_4d = op.Unsqueeze(cos, [1])  # convert
         sin_4d = op.Unsqueeze(sin, [1])
         return op.RotaryEmbedding(
@@ -226,7 +233,24 @@ def rewrite(
 _const_freqs = CosSinCacheFusion.rule("CosSinCache_const_freqs", cast=False, const_freqs=True)
 _basic = CosSinCacheFusion.rule("CosSinCache", cast=False)
 
-cos_sin_cache_rules = pattern.RewriteRuleSet([_cast, _cast_const_freqs, _const_freqs, _basic])
+cos_sin_cache_rules = pattern.RewriteRuleSet(
+    [
+        _cast,
+        _cast_const_freqs,
+        _const_freqs,
+        _basic,
+        *(
+            CosSinCacheFusion.rule(
+                f"CosSinCache_concat_after_trig_{cast}_{const_freqs}",
+                cast=cast,
+                const_freqs=const_freqs,
+                concat_after_trig=True,
+            )
+            for cast in (False, True)
+            for const_freqs in (False, True)
+        ),
+    ]
+)
 
 
 fuse_cos_sin_cache = _fusion_utils.apply_fusion_rules(cos_sin_cache_rules)