diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 418e375c1b..7e15c069f3 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/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, diff --git a/tests/function_libs/torch_lib/extra_opinfo.py b/tests/function_libs/torch_lib/extra_opinfo.py index d925be6877..2c97b73b22 100644 --- a/tests/function_libs/torch_lib/extra_opinfo.py +++ b/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, + ), ] diff --git a/tests/function_libs/torch_lib/ops_test_data.py b/tests/function_libs/torch_lib/ops_test_data.py index 241ac98cf9..2025c2ae27 100644 --- a/tests/function_libs/torch_lib/ops_test_data.py +++ b/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(