diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 99c20eca18..7a918cc029 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/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""" diff --git a/tests/function_libs/torch_lib/e2e_ops_tests.py b/tests/function_libs/torch_lib/e2e_ops_tests.py index 9e83f07625..df06d8294f 100644 --- a/tests/function_libs/torch_lib/e2e_ops_tests.py +++ b/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). diff --git a/tests/function_libs/torch_lib/extra_opinfo.py b/tests/function_libs/torch_lib/extra_opinfo.py index cea7f03b51..8caf341b8e 100644 --- a/tests/function_libs/torch_lib/extra_opinfo.py +++ b/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, diff --git a/tests/function_libs/torch_lib/ops_test_data.py b/tests/function_libs/torch_lib/ops_test_data.py index 3658a35b9d..ded5f634de 100644 --- a/tests/function_libs/torch_lib/ops_test_data.py +++ b/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",