TensorComprehensions
TensorComprehensions copied to clipboard
Add support for strided tensors
We currently incorrectly treat all tensors as dense. Need to add support for strided tensors. This can be done by adding the strides as extra params, and either constructing the tensor views in the emitted source using the strides instead of the sizes, or just doing the storage flattening ourselves before or during codegen.
checklist after talking to @ftynse [~@prigoyal~ @protonu is working on it]
- [ ] propagate DLTensor stride information to Halide (it has similar struct and already takes bits, lanes info, so should be easy)
- [ ] codegen, where we emit "tensor views" by casting a pointer into a multidimensional array - the casting should be aware of strides. Also, Halide->polyhedral, there is a place that creates constraints on tensor sizes for polyhedral; it may or may not be adapted to strides.
- [ ] add the stride info to the tuner proto
- [ ] make proper changes to the frontend (aten, pytorch, caffe2 etc.)
@abadams do you have time to look into this?
assigning this to myself, I'll work on it soon.
Sorry, snowed under with Halide work I neglected during the TC release.
no worries at all, I will probably have some time next week and can look into it :)