Optimize away zero-length concat operands (microsoft#2150)

We optimize `Concat (x1, x2, x3)` if one or more the concat operands has
zero length along the concatenated axis-dimension.

This pattern shows up, for example, in Phi models. See [this
line](https://github.com/huggingface/transformers/blob/786d9c5ed920a099573ea7b6dbf265f1aeb32fc0/src/transformers/models/phi3/modeling_phi3.py#L152)
in the implementation of partial-rotary-embedding:
```py
q_embed = torch.cat([(q_rot * cos) + (rotate_half(q_rot) * sin), q_pass], dim=-1)
```
In the special case of total-rotary-embedding, the second operand
`q_pass` of the concat is empty. This also interferes with the
pattern-matching for GQA in the generated graph. Optimizing the
redundant Concat away will help with GQA fusion as well.

Handle the edge case when all operands have zero size.

Author: gramalingam

onnxscript/optimizer/_constant_folding.py

@@ -558,21 +558,59 @@ def sequence_construct(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
 @register("Concat")
 def concat(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
     """Replace a Concat node with a single input by Identity"""
+
+    # Replace Concat(x) by Identity(x)
     inputs = node.inputs
     if len(inputs) == 1:
         return op.Identity(inputs[0])
-    # Track value of tensors that carry a shape value:
-    output = node.outputs[0]
-    if output is None:
+
+    axis = _get_int_attribute(node, "axis", None)
+    if axis is None:
         return None
+
+    # Eliminate zero-length operands from Concat
+    def has_zero_size(operand: ir.Value | None) -> bool:
+        if operand is None:
+            return False  # Invalid model
+        if (shape := operand.shape) is None:
+            return False
+        try:
+            # We have already checked that axis is an int value (!= None)
+            dim_size = shape[axis]  # type: ignore[index]
+        except IndexError:
+            return False
+        return dim_size == 0  # return False if symbolic or None or non-zero int value
+
+    new_inputs = [x for x in inputs if not has_zero_size(x)]
+    if len(new_inputs) != len(inputs):
+        if new_inputs:
+            # Remove zero-length operands from Concat
+            logger.debug(
+                "Concat: removing zero-length operand(s) %s => %s", inputs, new_inputs
+            )
+            return op.Concat(*new_inputs, axis=axis)
+        elif inputs:
+            # All operands are zero-length. Concat is a no-op, but we need to use one of the
+            # inputs to get the other dimensions correct:
+            logger.debug("Concat: removing all zero-length operands %s", inputs)
+            return op.Identity(inputs[0])
+        else:
+            # No inputs: invalid model.
+            return None
+
+    # Track value of tensors that carry a shape value:
+
     # Check axis attribute is 0
-    axis = _get_int_attribute(node, "axis", None)
+
     if axis != 0:
         return None
     shapes = [state.get_shape_value(input) for input in inputs]
     if any(shape is None for shape in shapes):
         return None
     concatenated = ir.Shape(dim for shape in shapes for dim in shape.dims)  # type: ignore[union-attr]
+    output = node.outputs[0]
+    if output is None:
+        return None
     state.set_sym_value(output, concatenated)
     return None
 

onnxscript/optimizer/_constant_folding_test.py

@@ -479,6 +479,44 @@ def test_concat_identity(self):
         self.assertEqual(len(optimized.graph), 1)
         self.assertEqual(optimized.graph.node(0).op_type, "Identity")
 
+    def test_concat_zero_length(self):
+        model = """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[N, 128] x1, float[N, 0] x2, float[N, 128] x3) => (float[N, M] z)
+            {
+                z = Concat <axis=-1> (x1, x2, x3)
+            }
+        """
+        optimized = self._fold(model)
+        self.assertEqual(len(optimized.graph), 1)
+        self.assertEqual([x.name for x in optimized.graph.node(0).inputs], ["x1", "x3"])
+
+    def test_concat_zero_length_identity(self):
+        model = """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[N, 0] x1, float[N, 128] x2, float[N, 0] x3) => (float[N, M] z)
+            {
+                z = Concat <axis=-1> (x1, x2, x3)
+            }
+        """
+        optimized = self._fold(model)
+        self.assertEqual(len(optimized.graph), 1)
+        self.assertEqual(optimized.graph.node(0).op_type, "Identity")
+        self.assertEqual([x.name for x in optimized.graph.node(0).inputs], ["x2"])
+
+    def test_concat_zero_length_output(self):
+        model = """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[N, 0] x1, float[N, 0] x2, float[N, 0] x3) => (float[N, M] z)
+            {
+                z = Concat <axis=-1> (x1, x2, x3)
+            }
+        """
+        optimized = self._fold(model)
+        self.assertEqual(len(optimized.graph), 1)
+        self.assertEqual(optimized.graph.node(0).op_type, "Identity")
+        self.assertEqual([x.name for x in optimized.graph.node(0).inputs], ["x1"])
+
     def test_expand_identity(self):
         model = """
             <ir_version: 7, opset_import: [ "" : 17]>