tvm
tvm copied to clipboard
apache
[Relax] Implement operators to read runtime DLTensor* information
Relax is capable of expressing tensors whose element type is unknown. However, these must typically be replaced with a known dtype prior to compilation, as most operators require known data types prior to legalization. This can be done by using a relax::MatchCast node, such as accepting a parameter arg: R.Tensor([16,16]), then defining the dtype using R.match_cast(arg, R.Tensor([16,16],'float16')).
However, using a R.match_cast node requires knowing which data type should be used in the new R.Tensor, and raises an error for an incorrect data type. If an argument may be one of two distinct data types, R.match_cast cannot be used to check which data type is in use.
This commit adds Relax operators to read the runtime values of a DLTensor* argument. These can be be used to normalize arguments prior to a compute step. For example, pre-processing a model weight that may be provided in either float16 or bfloat16 format.
The implementation seems solid, a very good change. I like the parameterized test cases too. The use of the _DLTensorShapeProxy resulted in an elegant UI in tvmscript.
Thank you! I really like polishing up the interface to be as clean as possible. (For my own sake if nothing else, as I am liable to forget a builtin-function name, but am much less likely to forget obj.dtype.)
Idle musing: I wonder if there's any way the PrimFuncs can be unrolled in cases where the parameters are known at compile time.
Prior to LegalizeOps, this is implicitly done by the FNormalize implementation for each operator. After LegalizeOps, not so much. While the FoldConstant pass can compile/run TIR functions if their arguments are known, there isn't a good way to indicate that a TIR function only requires the DLTensor struct, and not the data itself.
All CI tests passing, and thank you for the review @slyubomirsky ! I'll follow up with another PR to add inspection of the remainder of the DLTensor* fields.