Add ReverseSequence rule

scripts/optimize.py

@@ -2,6 +2,7 @@ import onnxscript
 import onnx_ir as ir
 import onnx_ir.passes.common
 import numpy as np
+import onnxslim
 
 
 class ReplaceDftWithMatMulRule(onnxscript.rewriter.RewriteRuleClassBase):
@@ -64,6 +65,37 @@ class RemoveCast(onnxscript.rewriter.RewriteRuleClassBase):
         return op.Identity(x)
 
 
+class RemoveReversedSequenceFork(onnxscript.rewriter.RewriteRuleClassBase):
+    def pattern(self, op, x, y, scale, bias):
+        x = op.Transpose(x)
+        y = op.Transpose(y)
+        x = op.ReverseSequence(x, _allow_other_inputs=True)
+        y = op.ReverseSequence(y, _allow_other_inputs=True)
+        x = op.Unsqueeze(x, _allow_other_inputs=True)
+        y = op.Unsqueeze(y, _allow_other_inputs=True)
+        concat = op.Concat(x, y)
+        mul = op.Mul(concat, scale)
+        add = op.Add(mul, bias)
+        return op.Transpose(add)
+
+    def rewrite(self, op, x, y, scale, bias, **kwargs):
+        # x: batch, 511, 96
+        neg_one = op.initializer(ir.tensor(np.array([-1], dtype=np.int64)), "neg_one")
+        int_64_min = op.initializer(
+            ir.tensor(np.array([-9223372036854775808], dtype=np.int64)), "int_64_min"
+        )
+        # slice
+        x = op.Slice(x, neg_one, int_64_min, neg_one, neg_one)
+        y = op.Slice(y, neg_one, int_64_min, neg_one, neg_one)
+        x = op.Unsqueeze(x, neg_one)
+        y = op.Unsqueeze(y, neg_one)
+        concat = op.Concat(x, y, axis=3)
+        # batch, 511, 96, 2
+        mul = op.Mul(concat, scale)
+        add = op.Add(mul, bias)
+        return op.Transpose(add, perm=[0, 3, 2, 1])  # batch, 2, 96, 511
+
+
 model = ir.load("model.onnx")
 
 # Set dynamic axes
@@ -72,7 +104,11 @@ model.graph.outputs[0].shape = ir.Shape(["batch", 6522])
 
 onnxscript.rewriter.rewrite(
     model,
-    [ReplaceDftWithMatMulRule().rule(), ReplaceSplit().rule(), RemoveCast().rule()],
+    [
+        ReplaceDftWithMatMulRule().rule(),
+        ReplaceSplit().rule(),
+        RemoveCast().rule(),
+    ],
 )
 
 # Change all int32 initializers to int64
@@ -131,6 +167,19 @@ onnxscript.optimizer.optimize(
     model, input_size_limit=1024 * 1024 * 1024, output_size_limit=1024 * 1024 * 1024
 )
 
+print("Slimming model...")
+model = ir.from_proto(onnxslim.slim(ir.to_proto(model)))
+
+print("Removing reversed sequence fork...")
+onnxscript.rewriter.rewrite(
+    model,
+    [
+        RemoveReversedSequenceFork.rule(),
+    ],
+)
+
+# Use onnxslim to do shape inference
+model = ir.from_proto(onnxslim.slim(ir.to_proto(model)))
 
 onnx_ir.passes.common.ClearMetadataAndDocStringPass()(model)
 model.graph.inputs[0].name = "input"