add cos, sin calc functions

Author: shubhambhokare1

onnxscript/rewriter/ort_fusions/cos_sin_cache.py

@@ -29,15 +29,11 @@
 #
 # This produces cos/sin values in a form that can be used by ORT's custom ops.
 
-# TODO: To apply the pattern-rewrite, we need to know the maximum position id.
-# Need to find a way to get this information from the model or its config.
-
 
 class CosSinCacheFusion(pattern.RewriteRuleClassBase):
     def __init__(
         self,
         name: str,
-        max_pos_id: int,
         *,
         cast: bool = False,
         reshape: bool = False,
@@ -47,13 +43,66 @@ def __init__(
         # matched nodes as part of the rewrite-step. We apply a separate final
         # pass to remove unused nodes.
         super().__init__(name, remove_nodes=False)
-        self._max_pos_id = max_pos_id
+        # TODO: Determine what should be the default max_pos_id value
+        self._max_pos_id = 2048  # Set a default max_pos_id value
+        self._max_pos_id_from_model = False
         # map from inv_freq to (cos, sin) values for transformed graph
         self._inv_freq_cos_sin_cache: dict[ir.Value, tuple[ir.Value, ir.Value]] = {}
         self._reshape = reshape
         self._cast = cast
         self._const_freqs = const_freqs
 
+    def _set_max_position_ids_from_model(self, model: ir.Model):
+        """Extract max_position_ids value from the metadata of an ONNX model."""
+        # TODO: Determine what the correct metadata key is for max_position_ids
+        if model.metadata_props["max_position_id"] is not None:
+            self._max_pos_id = model.metadata_props["max_position_id"]  # type: ignore[assignment]
+            self._max_pos_id_from_model = True
+
+    def _compute_const_freqs(self, op, freqs):
+        """Compute cos/sin values when frequencies are constant."""
+        angles = freqs.const_value.numpy()
+        cos_value = np.cos(angles)
+        sin_value = np.sin(angles)
+        cos_2d = op.Constant(value=ir.tensor(cos_value))
+        sin_2d = op.Constant(value=ir.tensor(sin_value))
+        return cos_2d, sin_2d
+
+    def _compute_dynamic_freqs(self, op, inv_freq, position_ids, dtype):
+        """Compute cos/sin values dynamically based on inv_freq and position_ids."""
+        if self._max_pos_id_from_model:
+            # Use max_pos_id from the model metadata
+            max_pos_id = self._max_pos_id
+        elif position_ids.const_value is not None:
+            # Calculate max_pos_id from the position_ids tensor
+            max_pos_id = int(np.max(position_ids.const_value.numpy()))
+            self._max_pos_id = max_pos_id
+        else:
+            # Dynamically compute max_pos_id from position_ids using ONNX ops
+            inv_freq = op.Reshape(inv_freq, op.Constant(value_ints=[1, -1]))
+            max_pos_id = op.ReduceMax(position_ids)
+            max_pos_id = op.Squeeze(max_pos_id, op.Constant(value_ints=[0]))
+            max_pos_id = op.Add(max_pos_id, op.Constant(value_int=1))
+            pos_id_range = op.Range(
+                op.Constant(value_int=0),
+                max_pos_id,
+                op.Constant(value_int=1),
+            )
+            pos_id_range = op.Reshape(pos_id_range, op.Constant(value_ints=[-1, 1]))
+            pos_id_range = op.Cast(pos_id_range, to=ir.DataType.FLOAT)
+            # Compute angles and cos/sin values
+            angles = op.MatMul(pos_id_range, inv_freq)
+            cos_2d = op.Cos(angles)
+            sin_2d = op.Sin(angles)
+            return cos_2d, sin_2d
+
+        # If we do not compute max_pos_id using ONNX ops, use inv_freq and position_ids
+        # to compute angles and cos/sin values
+        inv_freq_values = inv_freq.const_value.numpy().reshape(1, -1)
+        pos_id_range = np.arange(max_pos_id, dtype=np.float32).reshape(-1, 1)
+        angles = np.matmul(pos_id_range, inv_freq_values)
+        return self._compute_const_freqs(op, angles)
+
     def cleanup(self):
         self._inv_freq_cos_sin_cache.clear()
 
@@ -128,16 +177,11 @@ def rewrite(
         if inv_freq in self._inv_freq_cos_sin_cache:
             cos_2d, sin_2d = self._inv_freq_cos_sin_cache[inv_freq]
         else:
+            # Compute cos/sin values based on whether frequencies are constant
             if self._const_freqs:
-                angles = freqs.const_value.numpy()
+                cos_2d, sin_2d = self._compute_const_freqs(op, freqs)
             else:
-                inv_freq_values = inv_freq.const_value.numpy().reshape(1, -1)
-                pos_id_range = np.arange(self._max_pos_id, dtype=np.float32).reshape(-1, 1)
-                angles = np.matmul(pos_id_range, inv_freq_values)
-            cos_value = np.cos(angles)
-            sin_value = np.sin(angles)
-            cos_2d = op.Constant(value=ir.tensor(cos_value))
-            sin_2d = op.Constant(value=ir.tensor(sin_value))
+                cos_2d, sin_2d = self._compute_dynamic_freqs(op, inv_freq, position_ids, dtype)
             if self._cast:
                 cos_2d = op.Cast(cos_2d, to=dtype)
                 sin_2d = op.Cast(sin_2d, to=dtype)
@@ -157,15 +201,17 @@ def rewrite(
 
 
 _cast_const_freqs = CosSinCacheFusion.rule(
-    "CosSinCache_cast_const_freqs", 2048, cast=True, const_freqs=True
-)
-_cast = CosSinCacheFusion.rule("CosSinCache_cast", 2048, cast=True, const_freqs=False)
-_const_freqs = CosSinCacheFusion.rule(
-    "CosSinCache_const_freqs", 2048, cast=False, const_freqs=True
+    "CosSinCache_cast_const_freqs", cast=True, const_freqs=True
 )
-_basic = CosSinCacheFusion.rule("CosSinCache", 2048, cast=False)
+_cast = CosSinCacheFusion.rule("CosSinCache_cast", cast=True, const_freqs=False)
+_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])
 
 
-fuse_cos_sin_cache = _fusion_utils.apply_fusion_rules(cos_sin_cache_rules)
+def fuse_cos_sin_cache(model: ir.Model) -> int:
+    for rule in cos_sin_cache_rules:
+        if isinstance(rule, CosSinCacheFusion):
+            rule._set_max_position_ids_from_model(model)
+    return _fusion_utils.apply_fusion_rules(cos_sin_cache_rules)(model)