Fix: shape ops are never folded into quantized weights, leaving them mirrored vs dense

[kasper0406]

The bug

common::merge_affine_dequantize_with_consecutive_ops exists to fold shape-only ops (transpose, reshape, expand_dims, squeeze) into a constexpr op's quantized data, so they cost nothing at runtime. For the very common weight -> transpose -> matmul pattern it never fires, for two independent reasons — so a quantized weight ends up in the mirrored matmul layout relative to the identical dense weight.

1. The pass only matches the iOS16 op. The gate is if op.op_type != "constexpr_affine_dequantize": continue. Anything quantized against the iOS18 opset emits constexpr_blockwise_shift_scale, which never matches, so the pass is effectively dead for iOS18-quantized models.

2. Pass ordering pre-empts it. In PassPipeline.DEFAULT, fuse_transpose_matmul (index 53) folds the transpose into transpose_y long before the merge pass (84) can fold it into the weight. So even with (1) fixed, the transpose is already gone. Note _CLEANUP_PASSES orders these correctly — merge at 84, then fuse_transpose_matmul at 86 with the comment "there might be left over transpose … now can be fused back with matmul". The intended design is "fold into the weight first, fuse what's left"; the occurrence at 53 breaks it.

A dense weight is unaffected, because const_elimination (index 10) folds transpose(const) into a new const before either pass runs. const_elimination cannot do that for a constexpr op, by design. So today the same graph converts differently depending only on whether the weight is compressed.

Evidence

constexpr(data=[K,N], scale=[1,N]) -> transpose -> matmul(transpose_y=True) — the shape of every linear layer in a decoder-only LLM — after PassPipeline.DEFAULT:

                                  before                          after
blockwise_shift_scale (iOS18)   y=(32,8) transpose_y=False  ->  y=(8,32) transpose_y=True
affine_dequantize     (iOS16)   y=(32,8) transpose_y=False  ->  y=(8,32) transpose_y=True
dense fp16 const                y=(8,32) transpose_y=True   ->  y=(8,32) transpose_y=True

Both quantized paths were mirrored relative to dense; after the fix all three agree, and scale/offset are permuted along with the data. Repro script at the bottom.

This also has a real performance cost — a downstream user converting a 4-bit LLM measured decode going from 74.8 to 12.9 ms/token once the weights were forced into the dense layout (M4 Pro, macOS 26.6.2, .cpuAndGPU; I have not reproduced it). But the claim here is not that transpose_y=True is universally faster — only that a quantized weight should be laid out like the dense weight it replaces, and that folding a transpose into the weight is strictly better than folding it into a flag, since it removes the same op and additionally costs nothing at runtime.

The fix

1. Widen the merge pass to constexpr_blockwise_shift_scale. For this op scale/offset share data's rank with block_size[i] = data.shape[i] // scale.shape[i], so safety is provable per shape op: transpose, expand_dims and squeeze are always safe (apply the identical op to the parameters); reshape is excluded for blockwise and allowed only when parameters are a single element (per-tensor). This mirrors the existing SUPPORTED_OP_TYPES_PER_CHANNEL reasoning — a conservative subset, skipping anything not provably correct. offset gets the same treatment as scale; sub-byte dtypes survive; the pass declines when parameters are themselves produced by another constexpr op, which would undo the compression.

Deliberately out of scope: the lut variants. lut has rank K+2 with a vector_axis, a materially different correctness argument. Those keep today's behaviour exactly, so nothing regresses — clean follow-up.

2. Teach fuse_transpose_matmul to decline a transposed constexpr weight whose producer is an op type the merge pass supports, so the transpose survives to index 84.

This cannot lose a fusion: every pipeline containing _COMMON_PASSES also contains _CLEANUP_PASSES, which runs the merge pass and then fuse_transpose_matmul again — if the merge pass declines, the transpose is still fused at 86 by the exact fallback that comment describes. It is scoped to the supported op set, so constexpr_lut_to_dense and friends keep today's behaviour bit for bit, and non-constexpr operands are untouched.

Alternatives rejected: moving the merge pass earlier would run it before const_deduplication/dead_code_elimination, which it is annotated as depending on (it bails when a weight has multiple child ops), and would change folding behaviour ~30 passes earlier for all quantized models; running it in both positions has the same hazard plus a second full graph traversal per conversion.

Tests

New: 21 cases across TestBlockwiseShiftScaleConstElimination (20) and TestFuseTransposeMatmul (2, counting parametrization). 19 fail on unpatched main; the 2 that pass are negative tests that should pass either way. Coverage includes per-tensor / per-channel on either axis / true blockwise scales with and without offset, sub-byte dtype preservation, reshape folded only when per-tensor, the PassPipeline.DEFAULT ordering regression for both iOS16 and iOS18, and constexpr_lut_to_dense still being fused unchanged.

Correctness verified beyond op counts: every fold compared for exact equality against constexpr_blockwise_shift_scale.decompress of the original op; a standalone sweep over int4/int8, offsets, rank-4 permutations and op chains; and real Core ML predictions before vs after. Optimized vs unoptimized outputs differ by 2.3e-3 relative to output scale, and both sit the same distance from an fp32 numpy reference (2.2e-3 and 1.0e-3) — fp16 accumulation noise from the changed reduction order, not a regression.

Suite results: passes/tests/ 1973 passed, 6 failed (all 6 reproduce identically on unmodified main); ops/tests/iOS18/test_compression.py 2055 passed, 0 failed; test/optimize/coreml/ unchanged (644 pre-existing failures from missing sklearn/kmeans1d, byte-identical list before and after).

repro.py

"""
Repro for `merge_affine_dequantize_with_consecutive_ops` not folding shape ops into
iOS18-quantized weights, and for the `fuse_transpose_matmul` ordering interaction.

Pattern: a batch-1 GEMV against a per-output-channel weight, i.e. the shape of every
linear layer in a decoder-only LLM.

    constexpr(data=[K, N], scale=[1, N]) -> transpose(perm=[1, 0]) -> matmul(transpose_y=True)

Two things can happen to that transpose:
  (A) `fuse_transpose_matmul` folds it into the matmul flag
      => matmul(y=constexpr[K, N], transpose_y=False)
  (B) `merge_affine_dequantize_with_consecutive_ops` folds it into the weight bytes
      => matmul(y=constexpr[N, K] (scale [N, 1]), transpose_y=True)

(B) is what a *dense* weight gets, because `const_elimination` folds `transpose(const)` at
pipeline index 10, long before `fuse_transpose_matmul`.
"""

import numpy as np

import coremltools as ct
from coremltools.converters.mil.mil import Builder as mb
from coremltools.converters.mil.mil import types
from coremltools.converters.mil.mil.passes.pass_pipeline import (
    PassPipeline,
    PassPipelineManager,
)
from coremltools.converters.mil.testing_utils import (
    apply_pass_and_basic_check,
    get_op_types_in_program,
)

K, N = 32, 8
np.random.seed(0)


def show(title, prog):
    print(f"\n===== {title} =====")
    print("op types:", get_op_types_in_program(prog))
    for op in prog.functions["main"].operations:
        if op.op_type.startswith("constexpr_"):
            shapes = {n: (None if v is None else tuple(v.shape)) for n, v in op.inputs.items()}
            print(f"  {op.op_type}: {shapes}")
        if op.op_type == "matmul":
            print(
                f"  matmul: x={op.x.shape} y={op.y.shape} "
                f"transpose_x={None if op.transpose_x is None else op.transpose_x.val} "
                f"transpose_y={None if op.transpose_y is None else op.transpose_y.val}"
            )
        if op.op_type == "transpose":
            print(
                f"  transpose: perm={op.perm.val.tolist()} "
                f"in={op.x.shape} out={op.outputs[0].shape}"
            )


# ------------------------------------------------ root cause 1: the op-type gate
def root_cause_1():
    print("=== root cause 1: the merge pass only matches constexpr_affine_dequantize ===")
    data = np.random.randint(-8, 8, (4, 6)).astype(np.int8)

    @mb.program(input_specs=[], opset_version=ct.target.iOS18)
    def blockwise():
        w = mb.constexpr_blockwise_shift_scale(
            data=data, scale=np.random.rand(1, 6).astype(np.float16) + 0.1
        )
        return mb.transpose(x=w, perm=(1, 0))

    @mb.program(input_specs=[], opset_version=ct.target.iOS16)
    def affine():
        w = mb.constexpr_affine_dequantize(
            quantized_data=data,
            axis=1,
            scale=np.random.rand(6).astype(np.float32) + 0.1,
            zero_point=np.zeros(6, dtype=np.int8),
        )
        return mb.transpose(x=w, perm=(1, 0))

    for name, prog in (("blockwise_shift_scale", blockwise), ("affine_dequantize", affine)):
        before = get_op_types_in_program(prog)
        apply_pass_and_basic_check(prog, "common::merge_affine_dequantize_with_consecutive_ops")
        print(f"{name:24s} before={before}  after={get_op_types_in_program(prog)}")


# --------------------------------- root cause 2: fuse_transpose_matmul wins the race
def build_blockwise():
    data = np.random.randint(-8, 8, (K, N)).astype(np.int8)
    scale = np.random.rand(1, N).astype(np.float16) + 0.1

    @mb.program(
        input_specs=[mb.TensorSpec(shape=(1, K), dtype=types.fp16)],
        opset_version=ct.target.iOS18,
    )
    def prog(x):
        w = mb.constexpr_blockwise_shift_scale(data=data, scale=scale)
        return mb.matmul(x=x, y=mb.transpose(x=w, perm=(1, 0)), transpose_y=True)

    return prog


def build_affine():
    data = np.random.randint(-8, 8, (K, N)).astype(np.int8)

    @mb.program(input_specs=[mb.TensorSpec(shape=(1, K))], opset_version=ct.target.iOS16)
    def prog(x):
        w = mb.constexpr_affine_dequantize(
            quantized_data=data,
            axis=1,
            scale=np.random.rand(N).astype(np.float32) + 0.1,
            zero_point=np.zeros(N, dtype=np.int8),
        )
        return mb.matmul(x=x, y=mb.transpose(x=w, perm=(1, 0)), transpose_y=True)

    return prog


def build_dense():
    w_np = np.random.rand(K, N).astype(np.float16)

    @mb.program(
        input_specs=[mb.TensorSpec(shape=(1, K), dtype=types.fp16)],
        opset_version=ct.target.iOS18,
    )
    def prog(x):
        return mb.matmul(x=x, y=mb.transpose(x=mb.const(val=w_np), perm=(1, 0)), transpose_y=True)

    return prog


if __name__ == "__main__":
    print("coremltools", ct.__version__)
    root_cause_1()

    print("\n=== pipeline indices ===")
    for i, p in enumerate(PassPipeline.DEFAULT.passes):
        if p in (
            "common::const_elimination",
            "common::fuse_transpose_matmul",
            "common::merge_affine_dequantize_with_consecutive_ops",
        ):
            print(f"  {i:3d}  {p}")

    for name, builder in (
        ("blockwise_shift_scale (iOS18)", build_blockwise),
        ("affine_dequantize (iOS16)", build_affine),
        ("dense fp16 const (iOS18)", build_dense),
    ):
        prog = builder()
        show(f"{name}: BEFORE", prog)
        PassPipelineManager.apply_pipeline(prog, PassPipeline.DEFAULT)
        show(f"{name}: AFTER PassPipeline.DEFAULT", prog)