[Relax] Fix matmul and reductions with zero-size dimension return uninitialized memory

Author: cchung100m

python/tvm/relax/transform/legalize_ops/linear_algebra.py

@@ -45,6 +45,8 @@ def te_matmul(a: te.Tensor, b: te.Tensor) -> te.Tensor:
         b_relax = relax.Var("b", relax.TensorStructInfo(b.shape))
         f_infer_sinfo = call.op.get_attr("FInferStructInfo")
         output_shape = f_infer_sinfo(relax.op.matmul(a_relax, b_relax), bb).shape
+        if isinstance(a_shape[-1], trix.IntImm) and a_shape[-1] == 0:
+            return topi.full(output_shape, call.struct_info.dtype, 0)
 
         def matmul_compute(*idx_spatial):
             k = te.reduce_axis((0, a_shape[-1]), name="k")

python/tvm/relax/transform/legalize_ops/statistical.py

@@ -17,15 +17,54 @@
 # pylint: disable=invalid-name
 """Default legalization function for statistical operators."""
 
+from typing import Callable
+
 from tvm import te, tirx, topi
 
 from ...block_builder import BlockBuilder
 from ...expr import Call, Expr
 from .common import LegalizeFunc, TEFunc, register_legalize
 
 
-def _statistical(te_func: TEFunc) -> LegalizeFunc:
+def _normalize_reduction_axes(axis: list[int] | None, ndim: int) -> list[int]:
+    if axis is None:
+        return list(range(ndim))
+
+    axes = []
+    for dim in axis:
+        if isinstance(dim, trix.IntImm):
+            dim = dim.value
+        dim = int(dim)
+        axes.append(dim + ndim if dim < 0 else dim)
+    return axes
+
+
+def _has_const_zero_reduction_dim(call: Call) -> bool:
+    input_shape = call.args[0].struct_info.shape
+    if not isinstance(input_shape, ShapeExpr):
+        return false
+
+    axes = _normalize_reduction_axes(call.attrs.axis, len(input_shape.values))
+    return any(isinstance(input_shape.values[dim], tirx.IntImm) and input_shape.values[dim] == 0 for dim in axes)
+
+
+def _statistical(
+    te_func: TEFunc,
+    zero_dim_identity: int | float | bool | Callable[[str], int | float | bool] | None = None,
+) -> LegalizeFunc:
     def statistical_call_te(bb: BlockBuilder, call: Call) -> Expr:
+        if zero_dim_identity is not None and _has_const_zero_reduction_dim(call):
+            fill_value = (
+                zero_dim_identity(call.struct_info.dtype)
+                if callable(zero_dim_identity)
+                else zero_dim_identity
+            )
+            return bb.call_te(
+                topi.full,
+                call.struct_info.shape.values,
+                call.struct_info.dtype,
+                fill_value,
+            )
         return bb.call_te(te_func, call.args[0], call.attrs.axis, call.attrs.keepdims)
 
     return statistical_call_te
@@ -129,5 +168,8 @@ def _median(bb: BlockBuilder, call: Call) -> Expr:
 
 register_legalize("relax.max", _statistical(topi.max))
 register_legalize("relax.min", _statistical(topi.min))
-register_legalize("relax.prod", _statistical(topi.prod))
-register_legalize("relax.sum", _statistical(topi.sum))
+register_legalize(
+    "relax.prod",
+     _statistical(topi.prod, zero_dim_identity=lambda dtype: True if dtype == "bool" else 1),
+)
+register_legalize("relax.sum", _statistical(topi.sum, zero_dim_identity=0))