Skip to content

[TOPI] Use branchless boundary index for reflect/replicate pad #19928

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
junghyunpark2001 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

[TOPI] Use branchless boundary index for reflect/replicate pad #19928

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f298f5d
[TOPI] Use branchless boundary index for reflect/replicate pad
junghyunpark2001 Jul 2, 2026
e0db93f
[TOPI] Update replicate_pad autopad golden for branchless edge clamp
junghyunpark2001 Jul 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 10 additions & 19 deletions python/tvm/topi/nn/pad.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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))
Comment on lines +217 to +218

@gemini-code-assist gemini-code-assist Bot Jul 2, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

If size (which is data.shape[i]) and orig_idx have different integer types (e.g., int32 vs int64), the subtraction m - tvm.tirx.abs(orig_idx) can cause a type mismatch error during lowering. To ensure type safety and robustness, we should explicitly cast m to orig_idx.dtype using tvm.tirx.Cast.

Suggested change
m = size - 1
reflected_idx = m - tvm.tirx.abs(m - tvm.tirx.abs(orig_idx))
m = tvm.tirx.Cast(orig_idx.dtype, size - 1)
reflected_idx = m - tvm.tirx.abs(m - tvm.tirx.abs(orig_idx))

index_tuple.append(reflected_idx)
return data(*index_tuple)

Expand Down Expand Up @@ -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))

@gemini-code-assist gemini-code-assist Bot Jul 2, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Using a hardcoded "int32" for the constant 0 can lead to type mismatch errors if orig_idx is of type int64 (which is common when 64-bit indexing is enabled or for large tensors). To ensure robustness and prevent compilation failures, we should use orig_idx.dtype instead of "int32" and cast size - 1 to orig_idx.dtype as well.

Suggested change
clamped_idx = tvm.tirx.max(tvm.tirx.const(0, "int32"), tvm.tirx.min(size - 1, orig_idx))
clamped_idx = tvm.tirx.max(tvm.tirx.const(0, orig_idx.dtype), tvm.tirx.min(tvm.tirx.Cast(orig_idx.dtype, size - 1), orig_idx))

index_tuple.append(clamped_idx)
return data(*index_tuple)

Expand Down
Loading