From 8e000b2610774722103010abe1a79b9c9148d706 Mon Sep 17 00:00:00 2001 From: TANGBUDU <79953480+TANGBUDU@users.noreply.github.com> Date: Mon, 7 Sep 2026 09:58:32 +0200 Subject: [PATCH] Extend grouped_mm lowering to offset-based layouts Build on the existing converter from #2805. Unroll a static group count while preserving runtime offsets and logical output shapes. Add numerical, dynamic-input, empty-group, and compatibility tests. Assisted-by: ChatGPT --- .../function_libs/torch_lib/ops/core.py | 65 ++++- tests/function_libs/torch_lib/extra_opinfo.py | 42 +-- .../torch_lib/grouped_mm_test.py | 241 ++++++++++++++++++ .../function_libs/torch_lib/ops_test_data.py | 6 +- 4 files changed, 322 insertions(+), 32 deletions(-) create mode 100644 tests/function_libs/torch_lib/grouped_mm_test.py diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 27d30e377f..ce7c973ea4 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -4701,7 +4701,70 @@ def aten_grouped_mm( """_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.") + if self.shape is None or mat2.shape is None: + raise NotImplementedError("Grouped matmul requires known operand ranks.") + a_rank, b_rank = len(self.shape), len(mat2.shape) + if a_rank not in (2, 3) or b_rank not in (2, 3) or (a_rank == b_rank == 3): + raise ValueError("Grouped matmul with offsets requires at least one 2D operand.") + if offs.shape is None or len(offs.shape) != 1: + raise ValueError("Grouped matmul offsets must be 1D.") + groups = offs.shape[0] + if not isinstance(groups, int): + raise NotImplementedError( + "Grouped matmul requires a statically known number of groups." + ) + for shape in (self.shape, mat2.shape): + if len(shape) == 3 and isinstance(shape[0], int) and shape[0] != groups: + raise ValueError("Grouped matmul offsets and operand group counts must match.") + if bias is not None: + raise NotImplementedError("Grouped matmul with offsets does not support bias.") + if out_dtype is not None and out_dtype != -1 and out_dtype != self.dtype: + raise NotImplementedError( + "Grouped matmul with offsets requires the output dtype to match the input." + ) + + a_is_2d, b_is_2d = a_rank == 2, b_rank == 2 + if groups == 0: + output_shape = op.Concat( + op.Shape(self, start=a_rank - 2, end=a_rank - 1), + op.Shape(mat2, start=b_rank - 1, end=b_rank), + axis=0, + ) + if a_is_2d and b_is_2d: + output_shape = op.Concat(op.Constant(value_ints=[0]), output_shape, axis=0) + return op.Expand(op.CastLike(0, self), output_shape) + + # Only the group count is static. Boundaries remain values in the graph. + ends = op.Cast(offs, to=INT64.dtype) + start = op.Constant(value_ints=[0]) + outputs = [] + for i in range(groups): + end = op.Gather(ends, [i], axis=0) + a = ( + op.Slice(self, start, end, [1 if b_is_2d else 0]) + if a_is_2d + else op.Gather(self, i, axis=0) + ) + b = ( + op.Slice(mat2, start, end, [0 if a_is_2d else 1]) + if b_is_2d + else op.Gather(mat2, i, axis=0) + ) + result = op.MatMul(a, b) + if a_is_2d and b_is_2d: + result = op.Unsqueeze(result, [0]) + outputs.append(result) + start = end + + if a_is_2d and b_is_2d: + return op.Concat(*outputs, axis=0) + axis = 0 if a_is_2d else 1 + result = op.Concat(*outputs, axis=axis) + # PyTorch allocates the full output shape even if the last offset is short. + # Its unwritten tail is unspecified; fill it with zeros rather than shrink it. + length = op.Shape(self, start=0, end=1) if a_is_2d else op.Shape(mat2, start=1, end=2) + pads = op.Concat(op.Constant(value_ints=[0]), op.Sub(length, end), axis=0) + return op.Pad(result, pads, axes=[axis]) res = op.MatMul(self, mat2) if bias is not None: diff --git a/tests/function_libs/torch_lib/extra_opinfo.py b/tests/function_libs/torch_lib/extra_opinfo.py index c2fd220c8e..5f4b9aa486 100644 --- a/tests/function_libs/torch_lib/extra_opinfo.py +++ b/tests/function_libs/torch_lib/extra_opinfo.py @@ -28,38 +28,20 @@ def sample_inputs_grouped_mm(op_info, device, dtype, requires_grad, **kwargs): del kwargs make_arg = functools.partial( - torch_testing.make_tensor, - device=device, - dtype=dtype, - requires_grad=requires_grad, + torch_testing.make_tensor, device=device, dtype=dtype, requires_grad=requires_grad ) + # Native grouped_mm requires 16-byte-aligned strides, including on CPU. + for shape_a, shape_b in (((2, 3, 8), (2, 8, 8)), ((1, 2, 8), (1, 8, 8))): + yield opinfo_core.SampleInput(make_arg(shape_a), args=(make_arg(shape_b),)) - 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,), kwargs={"bias": bias_t}) - - # Test with bias and out_dtype - if dtype in (torch.float16, torch.bfloat16): - yield opinfo_core.SampleInput( - self_t, - args=(mat2_t,), - kwargs={"bias": bias_t, "out_dtype": torch.float32}, - ) + for shape_a, shape_b in ( + ((24, 8), (3, 8, 8)), + ((3, 8, 8), (8, 24)), + ((8, 24), (24, 8)), + ): + for boundaries in ([4, 12, 24], [8, 8, 24]): + offsets = torch.tensor(boundaries, dtype=torch.int32, device=device) + yield opinfo_core.SampleInput(make_arg(shape_a), args=(make_arg(shape_b), offsets)) def sample_inputs_scalar_tensor(op_info, device, dtype, requires_grad, **kwargs): diff --git a/tests/function_libs/torch_lib/grouped_mm_test.py b/tests/function_libs/torch_lib/grouped_mm_test.py new file mode 100644 index 0000000000..a1123bd2a9 --- /dev/null +++ b/tests/function_libs/torch_lib/grouped_mm_test.py @@ -0,0 +1,241 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +from __future__ import annotations + +import unittest + +import numpy as np +import onnx +import onnxruntime as ort +import parameterized +import torch +from torch.onnx._internal.exporter import _building, _tensors + +import onnxscript +from onnxscript import ir +from onnxscript.function_libs.torch_lib.ops import core + + +def _inputs(layout, dtype, length=24, groups=3): + shapes = { + "2d_3d": ((length, 8), (groups, 8, 8)), + "3d_2d": ((groups, 8, 8), (8, length)), + "2d_2d": ((8, length), (length, 8)), + } + rng = np.random.default_rng(0) + return tuple(rng.uniform(-1, 1, shape).astype(dtype) for shape in shapes[layout]) + + +def _reference(a, b, offsets): + # Fill only the groups PyTorch writes; unused trailing storage is unspecified. + if a.ndim == b.ndim == 2: + result = np.zeros((len(offsets), a.shape[0], b.shape[1]), dtype=a.dtype) + else: + result = np.zeros((a.shape[-2], b.shape[-1]), dtype=a.dtype) + start = 0 + for i, end in enumerate(offsets): + if a.ndim == b.ndim == 2: + result[i] = a[:, start:end] @ b[start:end] + elif a.ndim == 2: + result[start:end] = a[start:end] @ b[i] + else: + result[:, start:end] = a[i] @ b[:, start:end] + start = end + return result + + +def _capture(shapes, dtype, *, offset_shape, **kwargs): + opset = onnxscript.opset18 + specs = [("a", shapes[0], dtype), ("b", shapes[1], dtype)] + if offset_shape is not None: + specs.append(("offsets", offset_shape, ir.DataType.INT32)) + tensors = [ + _tensors.SymbolicTensor( + opset, name=name, shape=ir.Shape(shape), type=ir.TensorType(dt) + ) + for name, shape, dt in specs + ] + + tracer = _building.OpRecorder(opset, {}) + with onnxscript.evaluator.default_as(tracer): + result = core.aten_grouped_mm(*tensors, **kwargs) + rank = 3 if len(shapes[0]) == len(shapes[1]) else 2 + result.shape = ir.Shape([None] * rank) + result.dtype = kwargs.get("out_dtype", dtype) + graph = ir.Graph( + tensors, [result], nodes=tracer.nodes, opset_imports={"": 18}, name="grouped_mm" + ) + model = ir.to_proto(ir.Model(graph, ir_version=10)) + onnx.checker.check_model(model, full_check=True) + return model + + +def _session(model): + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + options.inter_op_num_threads = 1 + return ort.InferenceSession( + model.SerializeToString(), options, providers=["CPUExecutionProvider"] + ) + + +class GroupedMmTest(unittest.TestCase): + @parameterized.parameterized.expand( + [ + (f"{layout}_{dtype.__name__}_{case}", layout, dtype, offsets) + for layout in ("2d_3d", "3d_2d", "2d_2d") + for dtype in (np.float16, np.float32) + for case, offsets in ( + ("uneven", [4, 12, 24]), + ("empty_first", [0, 8, 24]), + ("empty_middle", [8, 8, 24]), + ("unused_tail", [0, 8, 16]), + ("all_empty", [0, 0, 0]), + ) + ] + ) + def test_offsets(self, _, layout, dtype, offsets): + a, b = _inputs(layout, dtype) + offsets = np.asarray(offsets, dtype=np.int32) + expected = _reference(a, b, offsets) + model = _capture( + (a.shape, b.shape), + ir.DataType.FLOAT16 if dtype == np.float16 else ir.DataType.FLOAT, + offset_shape=offsets.shape, + ) + actual = _session(model).run(None, {"a": a, "b": b, "offsets": offsets})[0] + self.assertEqual(actual.shape, expected.shape) + self.assertEqual(actual.dtype, expected.dtype) + rtol, atol = (1e-3, 1e-3) if dtype == np.float16 else (1e-5, 1e-6) + np.testing.assert_allclose(actual, expected, rtol=rtol, atol=atol) + if hasattr(torch.ops.aten, "_grouped_mm"): + native = torch.ops.aten._grouped_mm.default( + torch.from_numpy(a), torch.from_numpy(b), torch.from_numpy(offsets) + ).numpy() + self.assertEqual(actual.shape, native.shape) + # Do not compare uninitialized trailing rows/columns from the native op. + if layout == "2d_3d": + actual, native = actual[: offsets[-1]], native[: offsets[-1]] + elif layout == "3d_2d": + actual, native = actual[:, : offsets[-1]], native[:, : offsets[-1]] + np.testing.assert_allclose(actual, native, rtol=rtol, atol=atol) + + @parameterized.parameterized.expand([("2d_3d",), ("3d_2d",), ("2d_2d",)]) + def test_empty_group_list(self, layout): + a, b = _inputs(layout, np.float32, groups=0) + offsets = np.asarray([], dtype=np.int32) + model = _capture((a.shape, b.shape), ir.DataType.FLOAT, offset_shape=(0,)) + actual = _session(model).run(None, {"a": a, "b": b, "offsets": offsets})[0] + np.testing.assert_array_equal(actual, _reference(a, b, offsets)) + + def test_existing_dense_bias_and_cast(self): + a = np.ones((3, 8, 8), dtype=np.float32) + b = np.ones((3, 8, 8), dtype=np.float32) + model = _capture( + (a.shape, b.shape), + ir.DataType.FLOAT, + offset_shape=None, + bias=2.0, + out_dtype=ir.DataType.FLOAT16, + ) + actual = _session(model).run(None, {"a": a, "b": b})[0] + np.testing.assert_array_equal(actual, (a @ b + 2).astype(np.float16)) + + @parameterized.parameterized.expand([("2d_3d",), ("3d_2d",), ("2d_2d",)]) + def test_single_group_and_explicit_output_dtype(self, layout): + a, b = _inputs(layout, np.float32, groups=1) + offsets = np.asarray([24], dtype=np.int32) + model = _capture( + (a.shape, b.shape), + ir.DataType.FLOAT, + offset_shape=(1,), + out_dtype=ir.DataType.FLOAT, + ) + actual = _session(model).run(None, {"a": a, "b": b, "offsets": offsets})[0] + np.testing.assert_allclose(actual, _reference(a, b, offsets), rtol=1e-5, atol=1e-6) + + def test_dynamic_group_count_is_rejected(self): + with self.assertRaisesRegex(NotImplementedError, "statically known number of groups"): + _capture(((24, 8), ("groups", 8, 8)), ir.DataType.FLOAT, offset_shape=("groups",)) + + def test_operand_group_count_is_checked(self): + with self.assertRaisesRegex(ValueError, "group counts must match"): + _capture(((24, 8), (2, 8, 8)), ir.DataType.FLOAT, offset_shape=(3,)) + + def test_offset_rank_is_checked(self): + with self.assertRaisesRegex(ValueError, "1D"): + _capture(((24, 8), (3, 8, 8)), ir.DataType.FLOAT, offset_shape=(1, 3)) + + def test_two_dense_operands_with_offsets_are_rejected(self): + with self.assertRaisesRegex(ValueError, "2D operand"): + _capture(((3, 8, 8), (3, 8, 8)), ir.DataType.FLOAT, offset_shape=(3,)) + + def test_offset_bias_is_rejected(self): + with self.assertRaisesRegex(NotImplementedError, "bias"): + _capture(((24, 8), (3, 8, 8)), ir.DataType.FLOAT, offset_shape=(3,), bias=object()) + + def test_offset_output_dtype_change_is_rejected(self): + with self.assertRaisesRegex(NotImplementedError, "output dtype"): + _capture( + ((24, 8), (3, 8, 8)), + ir.DataType.FLOAT, + offset_shape=(3,), + out_dtype=ir.DataType.FLOAT16, + ) + + @parameterized.parameterized.expand([("2d_3d",), ("3d_2d",), ("2d_2d",)]) + def test_dynamic_offsets_and_shapes(self, layout): + shapes = { + "2d_3d": (("length", 8), (3, 8, 8)), + "3d_2d": ((3, 8, 8), (8, "length")), + "2d_2d": ((8, "length"), ("length", 8)), + }[layout] + model = _capture(shapes, ir.DataType.FLOAT16, offset_shape=(3,)) + session = _session(model) + self.assertIn("offsets", [value.name for value in session.get_inputs()]) + for size, values in ( + (24, [8, 16, 24]), + (24, [0, 8, 24]), + (24, [8, 8, 24]), + (24, [0, 8, 16]), + (32, [8, 16, 32]), + ): + with self.subTest(size=size, offsets=values): + a, b = _inputs(layout, np.float16, length=size) + offsets = np.asarray(values, dtype=np.int32) + actual = session.run(None, {"a": a, "b": b, "offsets": offsets})[0] + expected = _reference(a, b, offsets) + self.assertEqual(actual.shape, expected.shape) + np.testing.assert_allclose(actual, expected, rtol=1e-3, atol=1e-3) + + @parameterized.parameterized.expand([("2d_3d",), ("3d_2d",), ("2d_2d",)]) + @unittest.skipUnless(hasattr(torch.ops.aten, "_grouped_mm"), "requires aten::_grouped_mm") + def test_bfloat16_export(self, layout): + class Model(torch.nn.Module): + def forward(self, a, b, offsets): + return torch.ops.aten._grouped_mm.default(a, b, offsets) + + a, b = (torch.from_numpy(x).to(torch.bfloat16) for x in _inputs(layout, np.float32)) + offsets = torch.tensor([8, 16, 24], dtype=torch.int32) + length = 8 * torch.export.Dim("blocks", min=1) + dynamic_shapes = { + "2d_3d": ({0: length}, {}, {}), + "3d_2d": ({}, {1: length}, {}), + "2d_2d": ({1: length}, {0: length}, {}), + }[layout] + program = torch.onnx.export( + Model().eval(), + (a, b, offsets), + dynamo=True, + dynamic_shapes=dynamic_shapes, + optimize=False, + ) + onnx.checker.check_model(program.model_proto, full_check=True) + self.assertEqual( + program.model_proto.graph.output[0].type.tensor_type.elem_type, + onnx.TensorProto.BFLOAT16, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/function_libs/torch_lib/ops_test_data.py b/tests/function_libs/torch_lib/ops_test_data.py index 9c1f9088ac..8bb1625de4 100644 --- a/tests/function_libs/torch_lib/ops_test_data.py +++ b/tests/function_libs/torch_lib/ops_test_data.py @@ -743,7 +743,11 @@ 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).skip( + TorchLibOpInfo( + "ops.aten._grouped_mm", + core_ops.aten_grouped_mm, + tolerance={torch.float32: (1e-5, 1e-5)}, + ).skip( enabled_if=not hasattr(torch.ops.aten, "_grouped_mm"), reason="torch.ops.aten._grouped_mm is not available in this version of PyTorch", ),