[torchlib] Add aten::_grouped_mm ONNX lowering

onnxscript/function_libs/torch_lib/ops/core.py

@@ -99,6 +99,65 @@ def aten__log_softmax(self: TFloat, dim: int, half_to_float: bool) -> TFloatHigh
     return result
 
 
+@torch_op("aten::_grouped_mm", trace_only=True)
+def aten__grouped_mm(
+    self: TFloat,
+    mat2: TFloat,
+    offs: Optional[INT32] = None,
+    bias: Optional[TFloat] = None,
+    out_dtype: Optional[int] = None,
+) -> TFloat:
+    """_grouped_mm(Tensor self, Tensor mat2, Tensor? offs=None, Tensor? bias=None, ScalarType? out_dtype=None) -> Tensor"""
+
+    del out_dtype  # PyTorch currently requires the output dtype to match ``self``.
+    if bias is not None:
+        raise NotImplementedError("aten::_grouped_mm does not support bias in PyTorch")
+
+    self_is_2d = len(self.shape) == 2
+    mat2_is_2d = len(mat2.shape) == 2
+    if not self_is_2d and not mat2_is_2d:
+        if offs is not None:
+            raise ValueError("aten::_grouped_mm does not accept offsets for 3D operands")
+        return op.MatMul(self, mat2)
+
+    if offs is None:
+        raise ValueError("aten::_grouped_mm requires offsets when an operand is 2D")
+    group_count = offs.shape[0]
+    if not isinstance(group_count, int):
+        raise NotImplementedError(
+            "aten::_grouped_mm requires a statically known number of groups"
+        )
+    if group_count < 1:
+        raise ValueError("aten::_grouped_mm requires at least one group")
+
+    start = op.Constant(value_ints=[0])
+    outputs = []
+    for group_index in range(group_count):
+        end = op.Unsqueeze(op.Cast(op.Gather(offs, group_index, axis=0), to=INT64.dtype), [0])
+
+        if self_is_2d:
+            self_group = op.Slice(self, start, end, [1 if mat2_is_2d else 0])
+        else:
+            self_group = op.Gather(self, group_index, axis=0)
+
+        if mat2_is_2d:
+            mat2_group = op.Slice(mat2, start, end, [0 if self_is_2d else 1])
+        else:
+            mat2_group = op.Gather(mat2, group_index, axis=0)
+
+        output = op.MatMul(self_group, mat2_group)
+        if self_is_2d and mat2_is_2d:
+            output = op.Unsqueeze(output, [0])
+        outputs.append(output)
+        start = end
+
+    if self_is_2d and mat2_is_2d:
+        return op.Concat(*outputs, axis=0)
+    if self_is_2d:
+        return op.Concat(*outputs, axis=0)
+    return op.Concat(*outputs, axis=1)
+
+
 @torch_op("aten::_softmax", trace_only=True)
 def aten__softmax(self: TFloat, dim: int, half_to_float: bool) -> TFloatHighPrecision:
     """_softmax(Tensor self, int dim, bool half_to_float) -> Tensor"""

tests/function_libs/torch_lib/e2e_ops_tests.py

@@ -6,6 +6,7 @@
 import unittest
 
 import numpy as np
+import onnx
 import parameterized
 
 # TODO(pytorch/pytorch#129279): Migrate these tests to the PyTorch repo
@@ -97,6 +98,54 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         )
         _testing.assert_onnx_program(onnx_program)
 
+    @unittest.skipUnless(
+        hasattr(torch.ops.aten, "_grouped_mm"), "requires torch with aten::_grouped_mm"
+    )
+    def test_grouped_mm(self):
+        class Model(torch.nn.Module):
+            def forward(self, a, b):
+                return torch.nn.functional.grouped_mm(
+                    a.to(torch.bfloat16), b.transpose(-2, -1).to(torch.bfloat16)
+                ).to(torch.float32)
+
+        groups, rows, columns, contraction = 4, 16, 32, 64
+        a = torch.randn(groups, rows, contraction)
+        b = torch.randn(groups, columns, contraction)
+        onnx_program = torch.onnx.export(
+            Model().eval(),
+            (a, b),
+            dynamo=True,
+            dynamic_shapes=(
+                {0: "groups", 1: "rows", 2: "contraction"},
+                {0: "groups", 1: "columns", 2: "contraction"},
+            ),
+            optimize=False,
+        )
+
+        onnx.checker.check_model(onnx_program.model_proto, full_check=True)
+        self.assertIn("MatMul", [node.op_type for node in onnx_program.model_proto.graph.node])
+
+    @unittest.skipUnless(
+        hasattr(torch.ops.aten, "_grouped_mm"), "requires torch with aten::_grouped_mm"
+    )
+    def test_grouped_mm_with_offsets(self):
+        class Model(torch.nn.Module):
+            def forward(self, a, b, offsets):
+                return torch.ops.aten._grouped_mm.default(a, b, offsets)
+
+        offsets = torch.tensor([2, 5, 9], dtype=torch.int32)
+        a = torch.randn(9, 8, dtype=torch.bfloat16)
+        b = torch.randn(3, 8, 8, dtype=torch.bfloat16)
+        onnx_program = torch.onnx.export(
+            Model().eval(), (a, b, offsets), dynamo=True, optimize=False
+        )
+
+        onnx.checker.check_model(onnx_program.model_proto, full_check=True)
+        self.assertEqual(
+            [node.op_type for node in onnx_program.model_proto.graph.node].count("MatMul"),
+            3,
+        )
+
     def test_rand_like_memory_format(self):
         # These random *_like ops are non-deterministic, so assert the export
         # succeeds rather than comparing values (see issue #3002).

tests/function_libs/torch_lib/extra_opinfo.py

@@ -37,6 +37,44 @@ def sample_inputs_scalar_tensor(op_info, device, dtype, requires_grad, **kwargs)
         yield opinfo_core.SampleInput(item, dtype=dtype)
 
 
+def sample_inputs_grouped_mm(op_info, device, dtype, requires_grad, **kwargs):
+    """Sample all supported 2D/3D operand layouts for grouped matrix multiplication."""
+    del op_info
+    del kwargs
+
+    make_arg = functools.partial(
+        torch_testing.make_tensor, device=device, dtype=dtype, requires_grad=requires_grad
+    )
+    groups, rows, contraction, columns = 3, 5, 8, 8
+
+    # 3D x 3D: regular batched matrix multiplication.
+    yield opinfo_core.SampleInput(
+        make_arg((groups, rows, contraction)),
+        args=(make_arg((groups, contraction, columns)),),
+    )
+
+    # 2D x 3D: offsets split rows of the left operand.
+    row_offsets = torch.tensor([2, 5, 9], dtype=torch.int32, device=device)
+    yield opinfo_core.SampleInput(
+        make_arg((9, contraction)),
+        args=(make_arg((groups, contraction, columns)), row_offsets),
+    )
+
+    # 3D x 2D: offsets split columns of the right operand.
+    column_offsets = torch.tensor([4, 12, 16], dtype=torch.int32, device=device)
+    yield opinfo_core.SampleInput(
+        make_arg((groups, rows, contraction)),
+        args=(make_arg((contraction, 16)), column_offsets),
+    )
+
+    # 2D x 2D: offsets split the contraction dimension of both operands.
+    contraction_offsets = torch.tensor([4, 12, 16], dtype=torch.int32, device=device)
+    yield opinfo_core.SampleInput(
+        make_arg((rows, 16)),
+        args=(make_arg((16, columns)), contraction_offsets),
+    )
+
+
 def sample_inputs_linear_1d_weight(op_info, device, dtype, requires_grad, **kwargs):
     """Sample inputs for linear with a 1D weight (in_features,), no out_features dim.
 
@@ -2559,6 +2597,21 @@ def sample_inputs_masked_scatter(op_info, device, dtype, requires_grad, **kwargs
         sample_inputs_func=sample_inputs_bilinear,
         supports_out=False,
     ),
+    *(
+        (
+            opinfo_core.OpInfo(
+                "ops.aten._grouped_mm",
+                op=torch.ops.aten._grouped_mm.default,
+                aten_name="_grouped_mm",
+                dtypes=(torch.float32,),
+                sample_inputs_func=sample_inputs_grouped_mm,
+                supports_autograd=False,
+                supports_out=False,
+            ),
+        )
+        if hasattr(torch.ops.aten, "_grouped_mm")
+        else ()
+    ),
     opinfo_core.OpInfo(
         "ops.aten.linear.1d_weight",
         op=torch.nn.functional.linear,

tests/function_libs/torch_lib/ops_test_data.py

@@ -472,6 +472,11 @@ def _where_input_wrangler(
         fft_ops.aten__fft_r2c,
         tolerance={torch.float64: (2e-6, 2e-6), torch.float32: (3e-2, 3e-4)},
     ),
+    *(
+        (TorchLibOpInfo("ops.aten._grouped_mm", core_ops.aten__grouped_mm),)
+        if hasattr(torch.ops.aten, "_grouped_mm")
+        else ()
+    ),
     TorchLibOpInfo("ops.aten._local_scalar_dense", core_ops.aten__local_scalar_dense),
     TorchLibOpInfo(
         "ops.aten._log_softmax",