[Relax] Implement relax.op.view

[Lunderberg]

This commit implements relax.op.view (R.view in TVMScript) to produce a view into an existing array. This returned view shares the same backing allocation as the existing array.

Because R.view comes with potential trade-offs; such as increased memory footprint, performance cost to apply a non-zero DLTensor::byte_offset, and potential misalignment for vector operators; this PR does not use R.view apart from unit tests. Applications of R.view, either for specific compute kernels or in optimization passes, is instead kept for follow-up PRs.

[Lunderberg]

By line count, this PR looks bigger than it actually is. Because this functionality can , there's a lot more error checking than usual, and a lot more test cases to validate those error-checking paths. Breaking down the changes in this commit, the majority are test cases, and the majority of what remain are error-checking paths.

• ~100 lines implementation in view.cc
• ~260 lines error checking in view.cc
• ~100 lines exposing the functionality through the Python API
• ~750 lines testing the functionality in test_op_view.py

[masahi]

Why does this need to be a packed func? Can't we use the C++ function directly?

[Lunderberg]

If we want to define the StructInfo for the generated Call node, we could call the C++ function directly. Ideally, though, if the arguments change due to some downstream transform, the shape inference should be repeated with the new argument. For that

• Using an operator's FInferStructInfo function. This only applies to tvm.ir.Op instances, and we've just removed that as part of LegalizeOps.
• Using the params and ret fields from the FuncStructInfo. This works for static cases, and most dynamic shapes, but doesn't support inferring the output shape based on a ShapeExpr argument.
• Using FuncStructInfo::OpaqueFunc with a derivation func. This is the most general method, and essentially lets you pack a FInferStructInfo into a FuncStructInfo.

The third option doesn't come up very often, and I could only find one previous location where this functionality is used. (Previous usage is here, which defines the use of CallNode::sinfo_args as the default inferred struct info for external functions.)

[tqchen]

Just want to note here. Having view operation can in general cause problems mainly because most ops(including generated and external ones) assumes elem_offset = 0 for performance reasons(alignment, less kernel argument).

Ideally we would like to ensure such assumption to hold. This being said there are some needs in slicing out the arrays. e.g. in the case of LoRA elements. There are a few ways to go with this:

• Enable special ops that handle inputs which can come a view, likely only LoRA ones, which can inline view operations into ops.
• Add a R.memory.ensure_compact operation, which can potentially results in a copy for backends that do not have direct memory ptr access, but can potentially do ptr editing for backends that support them (this would need a target dependent lowering)
• R.view perhaps should be renamed as R.memory.view, this is a more advanced operator that contains certain assumptions and likely not something we want to advertise genrally for now.

[Lunderberg]

Having view operation can in general cause problems mainly because most ops(including generated and external ones) assumes elem_offset = 0

So long as the operations assert that the element offset is zero when this assumption is being made, this makes sense. This is what we do in MakePackedAPI for PrimFuncs that require the elem_offset to be zero.

For external operations that accept an aligned pointer to data, I like your suggestion of ensuring that we provide aligned data. My plan was to only introduce views that maintain the same alignment that is provided by existing allocations.

For external operations that accept a NDArray and assume the offset is zero without validating that assumption, I think we should view this as a bug in those external operations. If we know specific cases that ignore the offset, we can definitely insert alignment operators to provide aligned buffers. I wouldn't want to add it for every single external operation, though, because at some point we need to trust that functions accept the arguments that they say are accepted.

• Add a R.memory.ensure_compact operation, which can potentially results in a copy for backends that do not have direct memory ptr access, but can potentially do ptr editing for backends that support them (this would need a target dependent lowering)

To be clear, do you mean R.memory.ensure_aligned instead of R.memory.ensure_compact? This has come up a few times in this discussion, and I want to make sure that we are discussing the same thing. The view operation cannot be applied to a strided NDArray, and its output is always compact.

Having the dedicated operation for it would also work well for dynamically-shaped arguments. In those cases, we wouldn't know until runtime whether the operation requires a copy or not in order to provide an aligned argument.

Enable special ops that handle inputs which can come a view, likely only LoRA ones, which can inline view operations into ops.

I agree with aiming to have views be fused with later operations where possible, though I'd add that this is not LoRA-specific functionality. Anywhere that CombineParallelMatmul can be used to improve the matmul performance, a view into the result can be used to avoid an unnecessary copy from the output.

• R.view perhaps should be renamed as R.memory.view, this is a more advanced operator that contains certain assumptions and likely not something we want to advertise genrally for now.

I think the only assumption it makes is that a platform supports casting of pointers.

Regarding names, I agree that R.memory.view is a better name for it, and will update the PR.

[Lunderberg]

Changes have been made as requested, ready for re-review.

[tqchen]

R.memory.ensure_compact will actually return a new NDArray whose elem_offset=0, and for devices that allows explicit ptr moving, this is something that can accelerate backends

[Lunderberg]

Having R.memory.ensure_aligned makes sense, as that will allow us to explicitly mark where additional copies are required after applying R.memory.view, before passing to a compute kernel. I'm putting it together as a follow-up PR.

Any other concerns? I think this PR is ready to merge.

dismissing my previous requests as they are addressed

[tqchen]

Thank you! my previous comments are addressed. Unfortunately i didn't get a chance to do a thourough read, would be good to get review from another person. So i just dissmissed my previous comment

[Lunderberg]

Sounds good, and thank you on the feedback!

[masahi]

Need update here

[Lunderberg]

Sounds good, and updated.

[masahi]

Need update here

[Lunderberg]

Thank you on the catch, and updated.

[masahi]

Typo outputq?

And what is a derivation function?

[Lunderberg]

Typo, updated to output.

The derivation function is equivalent to FInferStructInfo set for relax operations, but can be applied to an arbitrary function. It can be used in cases where a PackedFunc may be called, and the output StructInfo should be derived from the input arguments, rather than taken from sinfo_args. This functionality has existed in the C++ API for a while, but it looks like this is the first time it's been exposed through the Python API, so I fleshed out the description here.

[masahi]

@tvm-bot rerun