diff --git a/python/tvm/topi/nn/pad.py b/python/tvm/topi/nn/pad.py index 8914c6ba485c..2319166ffbb7 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 _pad(*indices): 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 _pad(*indices): 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 024eb65d1965..b376ec39b4bd 100644 --- a/tests/python/relax/test_frontend_common.py +++ b/tests/python/relax/test_frontend_common.py @@ -119,32 +119,16 @@ def replicate_pad( 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