Update optimization script

scripts/optimize.py

@@ -1,15 +1,51 @@
 import onnxscript
 import onnx_ir as ir
 import onnx_ir.passes.common
+import numpy as np
 
 
 class ReplaceDftWithMatMulRule(onnxscript.rewriter.RewriteRuleClassBase):
-    
+    def pattern(self, op, x, dft_length):
+        x = op.Reshape(x, _allow_other_inputs=True)
+        dft = op.DFT(x, dft_length, _outputs=["dft_output"])
+        real_part = op.Slice(dft, [0], [1], [-1])
+        return op.Squeeze(real_part, [-1])
+
+    def rewrite(self, op, x: ir.Value, dft_length: ir.Value, dft_output: ir.Value):
+        # Get the DFT node attributes
+        dft_node = dft_output.producer()
+        assert dft_node is not None
+
+        dft_size = ir.convenience.get_const_tensor(dft_length).numpy().item()
+
+        # Create one-sided DFT matrix (only real part, DC to Nyquist)
+        # The real part of DFT is: Re(DFT[k]) = sum(x[n] * cos(2*pi*k*n/N))
+        # For one-sided DFT, we only need frequencies from 0 to Nyquist (dft_size//2 + 1)
+        num_freqs = dft_size // 2 + 1
+
+        # Vectorized creation of DFT matrix
+        k = np.arange(num_freqs, dtype=np.float32)[
+            :, np.newaxis
+        ]  # Shape: (num_freqs, 1)
+        n = np.arange(dft_size, dtype=np.float32)[np.newaxis, :]  # Shape: (1, dft_size)
+        dft_matrix = np.cos(
+            2 * np.pi * k * n / dft_size
+        )  # Shape: (num_freqs, dft_size)
+
+        # Create constant node for the DFT matrix
+        dft_matrix = op.initializer(ir.tensor(dft_matrix), name=f"{x.name}_dft_matrix")
+
+        # DFT axis is already at the end, direct matrix multiplication
+        result = op.MatMul(x, dft_matrix)
+
+        return result
+
 
 model = ir.load("model.onnx")
 onnxscript.optimizer.optimize(
     model, input_size_limit=1024 * 1024 * 1024, output_size_limit=1024 * 1024 * 1024
 )
+onnxscript.rewriter.rewrite(model, [ReplaceDftWithMatMulRule().rule()])
 
 onnx_ir.passes.common.ClearMetadataAndDocStringPass()(model)
 model.ir_version = 10