apache / apache/tvm

[Bug][Relax] Matmul rewrite passes drop explicit out_dtype and change float32 outputs to float16

Open
#20,198 0 comments 0 reactions 0 assignees View on GitHub
needs-triage type: bug
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.