[torchlib] Support offset-based grouped_mm layouts

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:

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):

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()

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",
     ),