This is an automated email from the ASF dual-hosted git repository.
tlopex 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 e4e489b788 [TOPI] Use branchless boundary index for reflect/replicate
pad (#19928)
e4e489b788 is described below
commit e4e489b788d08bf96fada95d221aa7c8a5fb566e
Author: junghyunpark <[email protected]>
AuthorDate: Fri Sep 11 04:17:18 2026 +0900
[TOPI] Use branchless boundary index for reflect/replicate pad (#19928)
reflect_pad and replicate_pad computed the boundary source index with
nested if_then_else, which lowers to per-element branches on CUDA and
made both modes ~1.4-1.5x slower than constant/circular since v0.21.0.
Replace them with branchless integer expressions that are bit-identical
over the valid pad domain:
reflect-101: m = size - 1; idx = m - abs(m - abs(orig_idx))
replicate: idx = max(0, min(size - 1, orig_idx))
Measured on RTX 4060 Ti (sm_89, CUDA 11.8): reflect 7.37->4.73us
(1.50x), replicate 7.14->4.31us (1.43x), output bit-identical. Also
removes the now-unused if_then_else import.
Fixes #19848
---------
Co-authored-by: Claude Opus 4.8 (1M context) <[email protected]>
---
python/tvm/topi/nn/pad.py | 29 ++++++++++-------------------
tests/python/relax/test_frontend_common.py | 28 ++++++----------------------
2 files changed, 16 insertions(+), 41 deletions(-)
diff --git a/python/tvm/topi/nn/pad.py b/python/tvm/topi/nn/pad.py
index 8914c6ba48..2319166ffb 100644
--- a/python/tvm/topi/nn/pad.py
+++ b/python/tvm/topi/nn/pad.py
@@ -18,7 +18,6 @@
import tvm
from tvm import te
-from tvm.tirx import if_then_else
from .. import tag
from ..utils import equal_const_int
@@ -211,15 +210,12 @@ def reflect_pad(data, pad_before, pad_after=None,
name="ReflectPadInput"):
orig_idx = idx - before
- reflected_idx = if_then_else(
- orig_idx < 0,
- -orig_idx, # reflect from start (no repeat)
- if_then_else(
- orig_idx >= size,
- (2 * size - 2) - orig_idx, # reflect from end
- orig_idx,
- ),
- )
+ # Branchless reflect-101 boundary index. This is bit-identical to
the
+ # nested if_then_else form (-orig_idx below 0, (2*size-2)-orig_idx
at or
+ # above size, orig_idx otherwise) over the valid reflect-pad
domain, but
+ # lowers to plain integer arithmetic instead of per-element
branches.
+ m = size - 1
+ reflected_idx = m - tvm.tirx.abs(m - tvm.tirx.abs(orig_idx))
index_tuple.append(reflected_idx)
return data(*index_tuple)
@@ -260,15 +256,10 @@ def replicate_pad(data, pad_before, pad_after=None,
name="ReplicatePadInput"):
before = pad_before[i]
orig_idx = idx - before
- clamped_idx = if_then_else(
- orig_idx < 0,
- tvm.tirx.const(0, "int32"), # replicate first element
- if_then_else(
- orig_idx >= size,
- size - 1, # replicate last element
- orig_idx,
- ),
- )
+ # Branchless edge clamp. This is bit-identical to the nested
+ # if_then_else form (0 below 0, size-1 at or above size, orig_idx
+ # otherwise) but lowers to min/max instead of per-element branches.
+ clamped_idx = tvm.tirx.max(tvm.tirx.const(0, "int32"),
tvm.tirx.min(size - 1, orig_idx))
index_tuple.append(clamped_idx)
return data(*index_tuple)
diff --git a/tests/python/relax/test_frontend_common.py
b/tests/python/relax/test_frontend_common.py
index 024eb65d19..b376ec39b4 100644
--- a/tests/python/relax/test_frontend_common.py
+++ b/tests/python/relax/test_frontend_common.py
@@ -119,32 +119,16 @@ class TestAutopad:
x[
T.int64(0),
T.int64(0),
- T.int64(0) : T.int64(4),
- T.int64(0) : T.int64(4),
+ T.max(T.int64(0), T.min(T.int64(3), v_i2)),
+ T.max(T.int64(0), T.min(T.int64(3), v_i3)),
]
)
T.writes(ReplicatePadInput[v_i0, v_i1, v_i2, v_i3])
ReplicatePadInput[v_i0, v_i1, v_i2, v_i3] = x[
- T.if_then_else(
- v_i0 < T.int64(0),
- T.int64(0),
- T.if_then_else(T.int64(1) <= v_i0, T.int64(0),
v_i0),
- ),
- T.if_then_else(
- v_i1 < T.int64(0),
- T.int64(0),
- T.if_then_else(T.int64(1) <= v_i1, T.int64(0),
v_i1),
- ),
- T.if_then_else(
- v_i2 < T.int64(0),
- T.int64(0),
- T.if_then_else(T.int64(4) <= v_i2, T.int64(3),
v_i2),
- ),
- T.if_then_else(
- v_i3 < T.int64(0),
- T.int64(0),
- T.if_then_else(T.int64(4) <= v_i3, T.int64(3),
v_i3),
- ),
+ T.int64(0),
+ T.int64(0),
+ T.max(T.int64(0), T.min(T.int64(3), v_i2)),
+ T.max(T.int64(0), T.min(T.int64(3), v_i3)),
]
@R.function
