This is an automated email from the ASF dual-hosted git repository.
MasterJH5574 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 277ae41efa [Relax] Legalize grouped conv with symbolic channels
(#20039)
277ae41efa is described below
commit 277ae41efa503f899359b3a581f7c54b4b87d197
Author: Guan-Ming Chiu <[email protected]>
AuthorDate: Fri Jul 24 03:31:55 2026 +0800
[Relax] Legalize grouped conv with symbolic channels (#20039)
- `LegalizeOps` skips grouped `conv1d/2d/3d` when channel size is
symbolic
- The only blocker is `topi.nn.conv`'s divisibility `assert`s, which
fail on symbolic `PrimExpr`; the grouped compute already handles
symbolic dims
## How
- Skip the divisibility check when the channel size is not a constant
int
- Remove the symbolic-channel guards from the conv legalize functions
---
python/tvm/relax/transform/legalize_ops/nn.py | 36 ----------------
python/tvm/topi/nn/conv2d.py | 6 ++-
.../python/relax/test_transform_legalize_ops_nn.py | 49 ++++++++++++++++++++++
3 files changed, 53 insertions(+), 38 deletions(-)
diff --git a/python/tvm/relax/transform/legalize_ops/nn.py
b/python/tvm/relax/transform/legalize_ops/nn.py
index a9c04d42e9..387c728902 100644
--- a/python/tvm/relax/transform/legalize_ops/nn.py
+++ b/python/tvm/relax/transform/legalize_ops/nn.py
@@ -42,18 +42,6 @@ def _nn_conv1d(bb: BlockBuilder, call: Call) -> Expr:
"cannot be legalized by TOPI at this moment."
)
return call
- if call.attrs.groups != 1:
- data_layout = s_tir.slayout(call.attrs.data_layout)
- kernel_layout = s_tir.slayout(call.attrs.kernel_layout)
- ic = call.args[0].ty.shape.values[data_layout.index_of("C")]
- oc = call.args[1].ty.shape.values[kernel_layout.index_of("O")]
- if not isinstance(ic, tirx.IntImm) or not isinstance(oc, tirx.IntImm):
- logging.info(
- "Conv1D where number of groups is more than one and input or
output "
- "channel size is symbolic cannot be legalized by TOPI at this
moment."
- )
- return call
-
return bb.call_te(
topi.nn.conv1d,
data=call.args[0],
@@ -83,18 +71,6 @@ def _nn_conv2d(bb: BlockBuilder, call: Call) -> Expr:
"cannot be legalized by TOPI at this moment."
)
return call
- if call.attrs.groups != 1:
- data_layout = s_tir.slayout(call.attrs.data_layout)
- kernel_layout = s_tir.slayout(call.attrs.kernel_layout)
- ic = call.args[0].ty.shape.values[data_layout.index_of("C")]
- oc = call.args[1].ty.shape.values[kernel_layout.index_of("O")]
- if not isinstance(ic, tirx.IntImm) or not isinstance(oc, tirx.IntImm):
- logging.info(
- "Conv2D where number of groups is more than one and input or
output "
- "channel size is symbolic cannot be legalized by TOPI at this
moment."
- )
- return call
-
return bb.call_te(
topi.nn.conv,
inp=call.args[0],
@@ -124,18 +100,6 @@ def _nn_conv3d(bb: BlockBuilder, call: Call) -> Expr:
"cannot be legalized by TOPI at this moment."
)
return call
- if call.attrs.groups != 1:
- data_layout = s_tir.slayout(call.attrs.data_layout)
- kernel_layout = s_tir.slayout(call.attrs.kernel_layout)
- ic = call.args[0].ty.shape.values[data_layout.index_of("C")]
- oc = call.args[1].ty.shape.values[kernel_layout.index_of("O")]
- if not isinstance(ic, tirx.IntImm) or not isinstance(oc, tirx.IntImm):
- logging.info(
- "Conv3D where number of groups is more than one and input or
output "
- "channel size is symbolic cannot be legalized by TOPI at this
moment."
- )
- return call
-
return bb.call_te(
topi.nn.conv,
inp=call.args[0],
diff --git a/python/tvm/topi/nn/conv2d.py b/python/tvm/topi/nn/conv2d.py
index 3268515b26..bdf6cbd24d 100644
--- a/python/tvm/topi/nn/conv2d.py
+++ b/python/tvm/topi/nn/conv2d.py
@@ -859,8 +859,10 @@ def conv(
if auto_scheduler_rewritten_layout:
raise RuntimeError("LEGACY-FLOW triggered, to be removed")
- assert in_channel % groups == 0, "input channels must divide group size"
- assert num_filter % groups == 0, "output channels must divide group size"
+ if isinstance(in_channel, int):
+ assert in_channel % groups == 0, "input channels must divide group
size"
+ if isinstance(num_filter, int):
+ assert num_filter % groups == 0, "output channels must divide group
size"
dilated_kernel_dimensions = [(k - 1) * dil + 1 for k, dil in
zip(kernel_dimensions, dilations)]
pad_begin, pad_end = get_pad_tuple_generic(padding,
dilated_kernel_dimensions)
diff --git a/tests/python/relax/test_transform_legalize_ops_nn.py
b/tests/python/relax/test_transform_legalize_ops_nn.py
index c70ea1b670..ab42dfadc6 100644
--- a/tests/python/relax/test_transform_legalize_ops_nn.py
+++ b/tests/python/relax/test_transform_legalize_ops_nn.py
@@ -435,6 +435,55 @@ def test_conv2d_symbolic():
tvm.ir.assert_structural_equal(mod, Expected)
+def test_conv2d_symbolic_group():
+ # fmt: off
+ @tvm.script.ir_module
+ class Conv2d:
+ @R.function
+ def main(x: R.Tensor(("n", "c", 28, 28), "float32"), w: R.Tensor(("f",
"c_div_8", 3, 3), "float32")) -> R.Tensor(("n", "f", 26, 26), "float32"):
+ n = T.int64()
+ f = T.int64()
+ gv: R.Tensor((n, f, 26, 26), "float32") = R.nn.conv2d(x, w,
groups=8)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("n", "c", 28, 28), dtype="float32"), w:
R.Tensor(("f", "c_div_8", 3, 3), dtype="float32")) -> R.Tensor(("n", "f", 26,
26), dtype="float32"):
+ n = T.int64()
+ f = T.int64()
+ gv = R.call_tir(Expected.conv2d, (x, w), out_ty=R.Tensor((n, f,
26, 26), dtype="float32"))
+ return gv
+
+ @T.prim_func(private=True, s_tir=True)
+ def conv2d(var_x: T.handle, var_w: T.handle, var_group_conv2d_nchw:
T.handle):
+ T.func_attr({"tirx.noalias": True})
+ n, c = T.int64(), T.int64()
+ x = T.match_buffer(var_x, (n, c, T.int64(28), T.int64(28)))
+ f, c_div_8 = T.int64(), T.int64()
+ w = T.match_buffer(var_w, (f, c_div_8, T.int64(3), T.int64(3)))
+ group_conv2d_nchw = T.match_buffer(var_group_conv2d_nchw, (n, f,
T.int64(26), T.int64(26)))
+ pad_temp = T.sblock_alloc_buffer((n, c, T.int64(28), T.int64(28)))
+ for i0, i1, i2, i3 in T.grid(n, c, T.int64(28), T.int64(28)):
+ with T.sblock("pad_temp"):
+ v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2,
i3])
+ T.reads(x[v_i0, v_i1, v_i2, v_i3])
+ T.writes(pad_temp[v_i0, v_i1, v_i2, v_i3])
+ pad_temp[v_i0, v_i1, v_i2, v_i3] = x[v_i0, v_i1, v_i2,
v_i3]
+ for nn, ff, yy, xx, rc, ry, rx in T.grid(n, f, T.int64(26),
T.int64(26), c // T.int64(8), T.int64(3), T.int64(3)):
+ with T.sblock("group_conv2d_nchw"):
+ v_nn, v_ff, v_yy, v_xx, v_rc, v_ry, v_rx =
T.axis.remap("SSSSRRR", [nn, ff, yy, xx, rc, ry, rx])
+ T.reads(pad_temp[v_nn, v_ff // (f // T.int64(8)) * (c //
T.int64(8)) + v_rc, v_yy + v_ry, v_xx + v_rx], w[v_ff, v_rc, v_ry, v_rx])
+ T.writes(group_conv2d_nchw[v_nn, v_ff, v_yy, v_xx])
+ with T.init():
+ group_conv2d_nchw[v_nn, v_ff, v_yy, v_xx] =
T.float32(0.0)
+ group_conv2d_nchw[v_nn, v_ff, v_yy, v_xx] =
group_conv2d_nchw[v_nn, v_ff, v_yy, v_xx] + pad_temp[v_nn, v_ff // (f //
T.int64(8)) * (c // T.int64(8)) + v_rc, v_yy + v_ry, v_xx + v_rx] * w[v_ff,
v_rc, v_ry, v_rx]
+ # fmt: on
+
+ mod = LegalizeOps()(Conv2d)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
def test_conv2d_transpose():
# fmt: off
@I.ir_module(s_tir=True)
