Skip to content

[Bug][Relax] AdjustMatmulOrder drops the outer matmul out_dtype #20202

Description

@lrcyyds1

Expected behavior

relax.transform.AdjustMatmulOrder() may reassociate a matmul chain, but it must preserve the original outer relax.matmul out_dtype. The function below should continue to return float32.

Actual behavior

For float16 inputs, AdjustMatmulOrder rebuilds both matmul calls with out_dtype="void". The function return dtype changes from float32 to float16.

This is not merely an expected floating-point reassociation difference: the observable Relax function signature and returned NDArray dtype change.

Observed locally:

before dtype: float32
after dtype:  float16
max_abs:      0.003734588623046875

The reverse reassociation direction has the same float32 to float16 regression.

### Environment

OS: Linux x86_64
Target: llvm
Relax VM exec_mode: bytecode
TVM commit: 5a8dae4d95c55c8fec9246a607a28c3ff54ffe05

### Steps to reproduce

import numpy as np
import tvm
from tvm import relax


def build_module(out_dtype):
    builder = relax.BlockBuilder()
    a = relax.Var("a", relax.TensorType((10, 2), "float16"))
    b = relax.Var("b", relax.TensorType((2, 20), "float16"))
    c = relax.Var("c", relax.TensorType((20, 2), "float16"))

    with builder.function("main", [a, b, c]):
        with builder.dataflow():
            lhs = builder.emit(relax.op.matmul(a, b))
            out = builder.emit(relax.op.matmul(lhs, c, out_dtype=out_dtype))
            output = builder.emit_output(out)
        builder.emit_func_output(output)

    return relax.transform.Normalize()(builder.finalize())


def run(mod, inputs):
    executable = relax.build(
        mod, target="llvm", relax_pipeline="default", exec_mode="bytecode"
    )
    vm = relax.VirtualMachine(executable, tvm.cpu())
    return vm["main"](*[tvm.runtime.tensor(value, tvm.cpu()) for value in inputs]).numpy()


rng = np.random.default_rng(20260827)
inputs = (
    rng.normal(0, 0.7, (10, 2)).astype("float16"),
    rng.normal(0, 0.7, (2, 20)).astype("float16"),
    rng.normal(0, 0.7, (20, 2)).astype("float16"),
)

before = build_module("float32")
after = relax.transform.AdjustMatmulOrder()(before)

expected = run(before, inputs)
actual = run(after, inputs)

print("before dtype:", expected.dtype)
print("after dtype:", actual.dtype)
print("max_abs:", np.max(np.abs(expected.astype("float32") - actual.astype("float32"))))

assert expected.dtype == np.dtype("float32")
assert actual.dtype == np.dtype("float16")

### Suspected cause

src/relax/transform/adjust_matmul_order.cc reconstructs the outer matmul with DataType::Void() in all rewrite branches. The inner new matmul should use DataType::Void(), but the new outer matmul should
preserve the original outer call's MatmulAttrs::out_dtype.

For example:

const auto* outer_call = expr.as<CallNode>();
const auto* outer_attrs = outer_call->attrs.as<MatmulAttrs>();
DataType outer_out_dtype = outer_attrs->out_dtype;

Then preserve outer_out_dtype on the reconstructed outer matmul.

### Scope

This occurs when AdjustMatmulOrder is explicitly included in a Relax pass pipeline. It is not triggered by the ordinary target-default LLVM pipeline alone.

### Triage

- needs-triage
- type: bug
- relax

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    needs-triagePRs or issues that need to be investigated by maintainers to find the right assignees to address ittype: bug

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions