tvm icon indicating copy to clipboard operation
tvm copied to clipboard

[TIR][Runtime] Allow use of external non-compact/non-aligned buffers

Open Lunderberg opened this issue 2 years ago • 5 comments
trafficstars

Prior to this commit, any use of tvm.nd.from_dlpack to create a strided NDArray, or a NDArray whose alignment was less than tvm::runtime::kAllocAlignment would raise an error. As a result, views into larger arrays, which are unlikely to be aligned and compact, could only be shared when copied into an aligned and compact buffer.

This commit moves the compact/aligned check from the NDArray class into the generated TIR code as part of DLTensor unpacking. These checks were initially introduced in https://github.com/apache/tvm/pull/11391, to avoid segfaults caused by use of non-aligned buffers in code intended for aligned buffers. The new checks will provide the same safeguard as the alignment is checked prior to use, but allows the alignment requirement to be relaxed on a per-buffer basis.

This approach also removes a potential bug resulting from compile-time configuration of tvm::runtime::kAllocAlignment, first introduced in https://github.com/apache/tvm/pull/13307. Since TVM supports cross-compiling, the installation of TVM used to compile a kernel may assume a larger value of kAllocAlignment than is provided by the runtime installation of TVM. By validating the alignment within the generated kernel, rather than as part of the runtime, this potential inconsistency would be caught.

Lunderberg avatar May 04 '23 21:05 Lunderberg

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

  • cc @Hzfengsy, @areusch, @junrushao, @quic-sanirudh, @shingjan See #10317 for details

Generated by tvm-bot

tvm-bot avatar May 04 '23 21:05 tvm-bot

It looks like this commit breaks most of the Rust tests, as the alignment assumed by the LLVM kernel was never validated when running through Rust. Unfortunately, I can't find a good way to produce aligned allocations in Rust, as Vec::new_in is unstable. For now, I've updated the Rust tests to opt out of requiring aligned buffers, but this isn't the best long-term solution.

Lunderberg avatar May 05 '23 14:05 Lunderberg

@areusch Could you take a look at the AOT change made in the most recent commit on this PR (link)? By changing the LLVM assumption to an assert, it's shaking out a number of locations where the assumption didn't hold, but I'm not very familiar with AOT.

Lunderberg avatar May 08 '23 17:05 Lunderberg

Rebased onto main to resolve conflict with https://github.com/apache/tvm/pull/15132. The individual commits are now cleaned up and ordered such that, if CI failures continue to occur, they can be segmented out into independent PRs.

Lunderberg avatar Jul 05 '23 18:07 Lunderberg

Rebased this onto main, as I attempt to clear out some of my (very) old PRs.

Lunderberg avatar May 23 '24 18:05 Lunderberg

Closing this PR, as it is unlikely that I will have a chance to get back to it.

If anybody else would like to pick this up, the last roadblock was that the Rust allocator does not provide the higher alignment requirement that is later assumed that memory allocation, and doesn't have a clear way to provide the larger-than-element-type alignment.

Lunderberg avatar Sep 11 '24 15:09 Lunderberg