This is not merely an expected floating-point reassociation difference: the observable Relax function signature and returned NDArray dtype change.
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
Expected behavior
relax.transform.AdjustMatmulOrder()may reassociate a matmul chain, but it must preserve the original outerrelax.matmulout_dtype. The function below should continue to returnfloat32.Actual behavior
For float16 inputs,
AdjustMatmulOrderrebuilds both matmul calls without_dtype="void". The function return dtype changes fromfloat32tofloat16.This is not merely an expected floating-point reassociation difference: the observable Relax function signature and returned NDArray dtype change.
Observed locally: