Skip to content

Commit 59f0a79

Browse files
committed
[WIP] torchdim Python port
Signed-off-by: Edward Yang <ezyang@meta.com> ghstack-source-id: 70c67dd Pull-Request: #160236
1 parent 482122d commit 59f0a79

20 files changed

+3738
-106
lines changed

.lintrunner.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,7 @@ is_formatter = true
123123
code = 'MYPY'
124124
include_patterns = [
125125
'setup.py',
126+
'functorch/dim/**/*.py',
126127
'torch/**/*.py',
127128
'torch/**/*.pyi',
128129
'caffe2/**/*.py',

functorch/dim/README.md

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -746,12 +746,14 @@ These compilers and language have syntax and semantics that resemble the loop-le
746746

747747
Dimension objects are just an extension of the existing PyTorch tensors and eager semantics, so there is no friction switching between normal Python code and code that uses them. However, since loops over the dimensions are defined implicitly, they can still execute in Python with good performance compared to explicit loops. Furthermore, with dimension objects, a tensors containing dimensions can compute through code that is oblivious to the dimension such as batching examples. There is no need to separate code into 'compiled' vs 'eager'.
748748

749-
In this way, first-class dims are a way of adapting the nicer syntax of these array compilers and languages to eager numpy-style libraries.
749+
In this way, first-class dims are a way of adapting the nicer syntax of these array compilers and languages to eager numpy-style libraries. Note, however, that first class dimensions are not natively compiled, so if you write code that performs many outer products with the expectation of it being fused, you will generally not get good performance or memory use (except for matrix-multiply-like patterns specifically.)
750750

751751

752752
Performance Expectations
753753
========================
754-
First-class dimensions are not a compiler. They provide syntax for existing PyTorch operations such as advanced indexing that is easier to read and write. For large sized tensors, the performance of any statements including them will be the same as using the already existing operations. An important exception is the pattern matching of products and summation, where performance will be improved by issuing to a matrix-multiply kernel. The C++ implementation of dimensions adds a small overhead of around 2us on top of PyTorch's normal overhead of 8us to each function that uses them. In the future, the implementation can incorporate more fusion optimization to further improve performance of this style of code.
754+
First-class dimensions are not a compiler. They provide syntax for existing PyTorch operations such as advanced indexing that is easier to read and write. For large sized tensors, the performance of any statements including them will be the same as using the already existing operations. An important exception is the pattern matching of products and summation, where performance will be improved by issuing to a matrix-multiply kernel.
755+
756+
Originally, there was a C++ implementation of dimensions adds a small overhead of around 2us on top of PyTorch's normal overhead of 8us to each function that uses them. However, this implementation had some manual memory managemetn bugs and was not kept up to date with CPython updates. The latest Python implementation is two orders of magnitude slower due to CPU overhead; for overhead sensitive applications you should compile the code to eliminate this overhead.
755757

756758

757759
## License

0 commit comments

Comments
 (0)