Add converter for aten::_grouped_mm.default

onnxscript/function_libs/torch_lib/ops/core.py

@@ -4510,6 +4510,29 @@ def aten_grid_sampler_3d_backward(
     raise NotImplementedError()
 
 
+@torch_op("aten::_grouped_mm", trace_only=True)
+def aten_grouped_mm(
+    self: TFloat,
+    mat2: TFloat,
+    offs: Optional[TInt] = None,
+    bias: Optional[TFloat] = None,
+    out_dtype: Optional[int] = None,
+) -> TFloat:
+    """_grouped_mm(Tensor self, Tensor mat2, *, Tensor? offs=None, Tensor? bias=None, int? out_dtype=None) -> Tensor"""
+
+    if offs is not None:
+        raise NotImplementedError(
+            "Grouped matmul with offsets (ragged/MoE) is not supported."
+        )
+
+    res = op.MatMul(self, mat2)
+    if bias is not None:
+        res = op.Add(res, bias)
+    if out_dtype is not None:
+        res = op.Cast(res, to=out_dtype)
+    return res
+
+
 def aten_gru_cell(
     input: TensorType,
     hx: TensorType,

tests/function_libs/torch_lib/extra_opinfo.py

@@ -23,6 +23,57 @@
 M = 10
 
 
+def sample_inputs_grouped_mm(op_info, device, dtype, requires_grad, **kwargs):
+    del op_info
+    del kwargs
+
+    make_arg = functools.partial(
+        torch_testing.make_tensor,
+        device=device,
+        dtype=dtype,
+        requires_grad=requires_grad,
+    )
+
+    cases = [
+        # (G, M, K), (G, K, N)
+        ((2, 3, 4), (2, 4, 5)),
+        ((1, 2, 2), (1, 2, 1)),
+    ]
+
+    for self_shape, mat2_shape in cases:
+        self_t = make_arg(self_shape)
+        mat2_t = make_arg(mat2_shape)
+
+        # Test without bias and without out_dtype
+        yield opinfo_core.SampleInput(self_t, args=(mat2_t,))
+
+        # Test with bias
+        g, _, _ = self_shape
+        _, _, n = mat2_shape
+        bias_t = make_arg((g, 1, n))
+        yield opinfo_core.SampleInput(self_t, args=(mat2_t, None, bias_t))
+
+        # Test with bias and out_dtype
+        if dtype in (torch.float16, torch.bfloat16):
+            yield opinfo_core.SampleInput(
+                self_t, args=(mat2_t, None, bias_t, torch.float32)
+            )
+
+
+def _mock_grouped_mm(self, mat2, offs=None, bias=None, out_dtype=None):
+    if hasattr(torch.ops.aten, "_grouped_mm"):
+        try:
+            return torch.ops.aten._grouped_mm(self, mat2, offs, bias, out_dtype)
+        except Exception:
+            pass
+    res = torch.matmul(self, mat2)
+    if bias is not None:
+        res = res + bias
+    if out_dtype is not None:
+        res = res.to(out_dtype)
+    return res
+
+
 def sample_inputs_scalar_tensor(op_info, device, dtype, requires_grad, **kwargs):
     del op_info
     del kwargs
@@ -3142,4 +3193,12 @@ def sample_inputs_masked_scatter(op_info, device, dtype, requires_grad, **kwargs
         sample_inputs_func=sample_inputs_masked_scatter,
         supports_out=False,
     ),
+    opinfo_core.OpInfo(
+        "ops.aten._grouped_mm",
+        aten_name="_grouped_mm",
+        op=_mock_grouped_mm,
+        dtypes=common_dtype.floating_types(),
+        sample_inputs_func=sample_inputs_grouped_mm,
+        supports_out=False,
+    ),
 ]

tests/function_libs/torch_lib/ops_test_data.py

@@ -750,6 +750,7 @@ def _where_input_wrangler(
         reason="fixme: ORT does not support empty tensors as input",
     ),
     TorchLibOpInfo("ge", core_ops.aten_ge),
+    TorchLibOpInfo("ops.aten._grouped_mm", core_ops.aten_grouped_mm),
     TorchLibOpInfo("gt", core_ops.aten_gt),
     TorchLibOpInfo("histc", core_ops.aten_histc)
     .skip(