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

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: 5a8dae4
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