[torchlib] Support offset-based grouped_mm layouts

[TANGBUDU]

Summary

Reworked on top of #2805 and current main (d1c005d1). This now extends the existing aten_grouped_mm converter rather than adding a second registration. The dense path, including its existing bias and cast handling, is unchanged.

The offsets branch supports:

• 2D x 3D: slice rows of the first operand, multiply by each group, concatenate rows.
• 3D x 2D: slice columns of the second operand, multiply by each group, concatenate columns.
• 2D x 2D: slice both contraction dimensions and stack the group results.

Only the number of groups must be static. Offset values and operand lengths remain runtime inputs. Tests cover uneven groups, empty groups, a single group, and an empty group list.

This is a follow-up to the dense support from #2805 and the fixed-group discussion in #2795, not a replacement for that merged PR.

Shape and compatibility details

For 2D x 3D and 3D x 2D, PyTorch allocates the full output extent even when the last offset does not cover the input. The earlier revision of this PR concatenated only the written prefix, incorrectly shortening the result. With inputs (24, 8) and (3, 8, 8) and offsets [0, 8, 16], that revision returned (16, 8) instead of (24, 8).

The new branch preserves the full extent and zero-fills the unwritten tail. PyTorch leaves that tail unspecified; numerical comparisons against the native operator check only the written portion, while checking the full output shape. The lowering assumes nondecreasing, in-bounds offsets.

The new offsets path rejects bias and output-dtype changes explicitly, matching the native operator tested here. It accepts an explicit output dtype equal to the input dtype. This does not remove the existing dense converter's bias/cast behavior: a separate direct-lowering test covers it.

Existing grouped-mm OpInfo inputs had strides rejected by native PyTorch. They now use aligned shapes and include all three offset layouts. Native bias samples were replaced because PyTorch rejects them; dense ONNX bias/cast coverage remains in the new tests. Float32 OpInfo comparisons use rtol=atol=1e-5 for MatMul accumulation differences.

Why static unrolling

This uses ordinary tensor operators and preserves BF16 graph types without introducing sequences or control-flow subgraphs. Graph size grows with the group count; dynamic group counts are explicitly unsupported.

I also tested a separate 2D x 3D SplitToSequence / SequenceMap prototype with both 3 and 5 groups in the same ONNX Runtime session. That approach works and produces a smaller graph, but SequenceMap's current type constraints do not include BF16. A dynamic-group implementation deserves its own dtype and subgraph coverage rather than being mixed into this change. No runtime performance advantage is claimed for unrolling.

Validation

Command used for both environments:

OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 \
python -m pytest \
  tests/function_libs/torch_lib/grouped_mm_test.py \
  tests/function_libs/torch_lib/ops_test.py \
  -k grouped_mm -q

• Python 3.12, PyTorch 2.14.0+cpu, ONNX 1.22.0, ONNX Runtime 1.29.0: 50 passed, 1 skipped; 23 subtests passed.
• Repository-pinned PyTorch 2.7.1+cpu, ONNX 1.18.0, ONNX Runtime 1.23.0, ONNX IR 0.1.16: 46 passed, 5 skipped; 15 subtests passed. Native-operator/export tests skip when _grouped_mm is unavailable.
• Restoring only the upstream converter makes all 30 new numerical offset cases fail with the expected NotImplementedError. Restoring this patch makes them pass.
• Executed FP16/FP32 ONNX graphs are compared against an independent per-group reference and, where available, native PyTorch. The dynamic tests reuse each graph with changed offsets and input lengths.
• Three real BF16 torch.onnx.export cases with dynamic shapes pass the ONNX checker and retain BF16 output types. These are export/checker tests, not BF16 runtime or GPU tests.
• Adjacent matrix tests (-k '__mm_ or __bmm_ or __addmm_ or __mv_') have the same four FP16 addmm subtest failures on this branch and a clean upstream worktree; failure names and messages match exactly. No additional failures in that selection.
• lintrunner -a and git diff --check: clean.

The full e2e file did not complete within the validation timeout. No full-suite pass, GPU validation, or performance improvement is claimed.

AI assistance

AI assistance was used for source review, implementation, comparison experiments, and test execution. The validation scope and limitations are recorded above.

[TANGBUDU]

@microsoft-github-policy-service agree