tlopex commented on issue #19848: URL: https://github.com/apache/tvm/issues/19848#issuecomment-4813136459
@junghyunpark2001 Thanks for the detailed write-up and benchmarks — this is a clear and well-scoped improvement. A PR is very welcome! The root cause is exactly as you describe: reflect_pad and replicate_pad compute the boundary index with a nested if_then_else per output element, whereas constant and circular use a single expression. The nested branches turn into extra comparison/select instructions on CUDA, which is what makes those two modes ~1.5x slower. The branchless rewrites are equivalent integer arithmetic, not an approximation, so the output stays bit-identical: - reflect-101: (size-1) - |(size-1) - |orig_idx|| - replicate (edge): clamp(orig_idx, 0, size-1) I verified this independently by exhaustively comparing the two index formulas: replicate matches the current logic for every shape/pad combination, and reflect matches for every well-defined configuration (pad <= size-1). The only divergence is in the over-pad region (pad > size-1), where reflect-101 is mathematically undefined and the current code already produces an out-of-bounds index — so that case is pre-existing and out of scope here. A couple of small notes for the PR: - Since pad.py already does from tvm import te, no new import is needed for te.abs/te.min/te.max. After the change, the top-level from tvm.tirx import if_then_else may become unused (the other functions use the fully-qualified tvm.tirx.if_then_else), so it can likely be dropped - mirror_pad has a similar nested-if_then_else pattern, but its semantics differ and it's a separate path — fine to leave for a follow-up Go ahead and open the PR — I'll help review. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
