tvm icon indicating copy to clipboard operation
tvm copied to clipboard

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

Open vacu9708 opened this issue 8 months ago • 1 comments

Summary

I was trying to resolve https://github.com/apache/tvm/issues/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 avatar Jun 14 '25 13:06 vacu9708

cc @tlopex

Hzfengsy avatar Jun 16 '25 01:06 Hzfengsy

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

Hzfengsy avatar Jul 07 '25 05:07 Hzfengsy

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

vacu9708 avatar Jul 08 '25 11:07 vacu9708