[Bug][Relax] Matmul rewrite passes drop explicit out_dtype and change float32 outputs to float16
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 4k
- Avg merge
- 2d 1h
- Merged PRs (30d)
- 112
Description
### Expected behavior
Relax matmul rewrite passes should preserve the original R.matmul out_dtype.
For an input program using float16 operands and:
R.matmul(x, weight, out_dtype="float32")
the transformed program should still return float32.
### Actual behavior
Two Relax rewrite passes appear to drop the original matmul out_dtype when creating replacement matmul calls:
ExpandMatmulOfSum
ReorderTakeAfterMatmul
For float16 operands with out_dtype="float32", the original function returns float32. After either rewrite pass, the transformed function returns float16.
Observed locally:
ExpandMatmulOfSum: return dtype float32 -> float16
ReorderTakeAfterMatmul: return dtype float32 -> float16
The transformed IR contains replacement matmuls with:
out_dtype="void"
This changes the observable output dtype and numerical behavior.
### Environment
OS: Linux x86_64
Python: 3.10.12
TVM version: 0.26.dev1
TVM commit: 5a8dae4d95c55c8fec9246a607a28c3ff54ffe05
Target: llvm
Relax VM exec_mode: compiled
### Steps to reproduce
import numpy as np
import tvm
from tvm import relax
def make_var(name, shape):
return relax.Var(name, relax.TensorType(shape, "float16"))
def build_expand_case():
bb = relax.BlockBuilder()
x = make_var("x", (2, 257))
a = make_var("A", (257, 31))
b = make_var("B", (257, 31))
with bb.function("main", params=[x, a, b]):
with bb.dataflow():
weight_sum = bb.emit(relax.op.add(a, b))
out = bb.emit(relax.op.matmul(x, weight_sum, out_dtype="float32"))
gv = bb.emit_output(out)
bb.emit_func_output(gv)
return bb.get()
def build_take_case():
bb = relax.BlockBuilder()
x = make_var("x", (2, 257))
weight_table = make_var("weight_table", (257, 47))
indices = relax.const(list(range(1, 47, 2)), "int64")
with bb.function("main", params=[x, weight_table]):
with bb.dataflow():
weight = bb.emit(relax.op.take(weight_table, indices, axis=1))
out = bb.emit(relax.op.matmul(x, weight, out_dtype="float32"))
gv = bb.emit_output(out)
bb.emit_func_output(gv)
return bb.get()
def return_dtype(mod):
text = mod.script(show_meta=False)
header = text.split("-> R.Tensor(", 1)[1].split(":", 1)[0]
return header.split('dtype="', 1)[1].split('"', 1)[0]
for name, mod, transform in [
("ExpandMatmulOfSum", build_expand_case(), relax.transform.ExpandMatmulOfSum()),
("ReorderTakeAfterMatmul", build_take_case(), relax.transform.ReorderTakeAfterMatmul()),
]:
before = relax.transform.Normalize()(mod)
after = transform(before)
print(name)
print("return dtype:", return_dtype(before), "->", return_dtype(after))
print(after.script(show_meta=False))
Observed output:
ExpandMatmulOfSum
return dtype: float32 -> float16
ReorderTakeAfterMatmul
return dtype: float32 -> float16
For ExpandMatmulOfSum, the original IR contains:
out: R.Tensor((2, 31), dtype="float32") = R.matmul(x, weight_sum, out_dtype="float32")
The transformed IR contains:
lv: R.Tensor((2, 31), dtype="float16") = R.matmul(x, A, out_dtype="void")
lv1: R.Tensor((2, 31), dtype="float16") = R.matmul(x, B, out_dtype="void")
gv: R.Tensor((2, 31), dtype="float16") = R.add(lv, lv1)
For ReorderTakeAfterMatmul, the original IR contains:
out: R.Tensor((2, 23), dtype="float32") = R.matmul(x, weight, out_dtype="float32")
The transformed IR contains:
lv: R.Tensor((2, 47), dtype="float16") = R.matmul(x, weight_table, out_dtype="void")
gv: R.Tensor((2, 23), dtype="float16") = R.take(lv, ..., axis=1, mode="fast")
I also ran an execution oracle comparing transformed output against NumPy reference semantics for the original IR. The transformed functions return float16 and differ numerically from the original float32
result.
A local patch oracle that preserves the original out_dtype in replacement matmuls restores the return dtype in all tested cases.
This is about public Relax rewrite passes invoked explicitly. I have not observed these passes being triggered automatically by the current CPU/GPU default target pipeline.
### Triage
- needs-triage
- type: bug
- relax
- transform
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with the Relax transform entry points ExpandMatmulOfSum and ReorderTakeAfterMatmul, using the reproduction programs in this issue to inspect their replacement matmul calls. Verify that explicit out_dtype="float32" is preserved and add or run coverage for both float16-operand cases. Done means transformed results retain float32 output and match the original NumPy reference semantics.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 72/100