Skip to content

Conversation

@vacu9708
Copy link
Contributor

@vacu9708 vacu9708 commented Jun 14, 2025

Summary

I was trying to resolve #18004, where an ONNX model causes a segmentation fault in TVM but not in onnxruntime.
Why the seg fault occurs
This occurs because take()(used in Compress)defaults to fast mode, which deliberately segfaults on out-of-bounds indices. (I guess for the sake of fast speed)
Solution
onnxruntime's Compress ignores out-of-bounds indices, so I initially considered matching that behavior in TVM. However, the ONNX specification doesn’t actually mandate how to handle out-of-bounds indices
Question for maintainers
Would you prefer modifying Compress to slice out invalid indices in order to mirror onnxruntime? although I guess this is unnecessary.

Instead of resolving this seg-fault, for now, this PR focuses just on enhancing take().

Updates

  • Add a mode parameter totake() in Relax
    • There are already several modes implemented in the transform layer, but Relax’s take() defaults to fast with no choice. I've exposed all modes so future use cases can pick as needed down the road.

  • Add NaN mode to take()
    • I was initially adding a raise mode(default of numpy.take()), but since there are no precedents raising runtime errors in the transform layer, I added an NaN mode that returns NaN for any out-of-bounds index. I'll remove this if this is unnecessary.

  • Add unit tests covering all take() modes
    • The tests for take() modes had been lost during a refactor.

  • Add a warning log for fast mode along an axis
    • take() along a flattened input shows a warning that lets users know about the potential seg fault, while take() along an axis does not. I think both take() need to warn about it.

  • Unify default modes in lower layers to fast for consistency with Relax
    • The default mode for take() in Relax is currently fast, I’ve made the lower layers use fast mode as well for consistency.

@vacu9708 vacu9708 force-pushed the compress branch 7 times, most recently from 8d27973 to fd051de Compare June 14, 2025 16:05
@Hzfengsy
Copy link
Member

cc @tlopex

@Hzfengsy
Copy link
Member

Hzfengsy commented Jul 7, 2025

Sorry for the late response. Please rebase the code, and we can get it in

@vacu9708 vacu9708 force-pushed the compress branch 13 times, most recently from fda34d2 to b0f02bd Compare July 8, 2025 00:36
- Add a `mode` parameter to Relax’s `take()`
- Add `NaN` mode to `take()`
- Add unit tests covering all `take()` modes
- Add a warning log for `fast` mode
- Unify default modes in lower layers to `fast` for consistency with Relax
@vacu9708
Copy link
Contributor Author

vacu9708 commented Jul 8, 2025

@Hzfengsy I've rebased the code according to the newly introduced reflection mechanism.

@Hzfengsy Hzfengsy merged commit 79368ce into apache:main Jul 8, 2025
12 checks passed
@vacu9708 vacu9708 deleted the compress branch July 8, 2025 13:07
ShiboXing pushed a commit to ShiboXing/tvm that referenced this pull request Aug 10, 2025
…ake() (apache#18061)

[Relax][Transform] Add mode choice, NaN mode, and warning for take()

- Add a `mode` parameter to Relax’s `take()`
- Add `NaN` mode to `take()`
- Add unit tests covering all `take()` modes
- Add a warning log for `fast` mode
- Unify default modes in lower layers to `fast` for consistency with Relax
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] TVM FFI encountered a Segfault when calling the invoke_stateful function

2 participants