Prevent JIT from overspecializing to every single size configuration by apaszke · Pull Request #10844 · pytorch/pytorch

apaszke · 2018-08-24T01:45:06Z

Please review the expects carefully to make sure there are no regressions. I tried to go over them one by one when they changed, but it's sometimes easy to miss finer details.

Summary of changes:

Renamed TensorType to CompleteTensorType. Added a new TensorType which records only the scalar type, number of dimensions, and device of a value. The argument behind the rename is to encourage people to use CompleteTensorType less, as most passes will only have limited information available. To make transition easier complete_type->cast<TensorType>() works, and makes our passes work with both kinds of specialization if they don't need extra the extra detail.
Renamed ArgumentSpec to CompleteArgumentSpec. Added a new ArgumentSpec, which matches argument only at the level of the new TensorType.
Shape analysis can process graphs with both CompleteTensorType and TensorType.
Fuser was a part that heavily relied on full shape information being available. Now, we simply try to fuse the largest possible graphs, and have to do run-time checks to make sure they match the code we generate. If they don't, we fall back to regular interpretation. The shape checks are implementing using an optimized method exploiting algebraic properties of shapes with broadcasting, and the relations of broadcasting with pointwise ops. A full written proof of correctness of the shape checking algorithm is included in a comment in graph_fuser.cpp.

@zdevito @ezyang @mruberry @ngimel @csarofeen

test/expect/TestJit.test_broadcast_fusion_cuda.expect

torch/csrc/jit/fusion_compiler.cpp

torch/csrc/jit/type.h

torch/csrc/jit/passes/shape_analysis.cpp

torch/csrc/jit/type.h

torch/csrc/jit/fusion_compiler.cpp

torch/csrc/jit/fusion_compiler.h

torch/csrc/jit/fusion_compiler.cpp

torch/csrc/jit/passes/graph_fuser.cpp

ezyang · 2018-08-24T15:26:14Z

Question: Shouldn't TensorType (the non-complete version) also track contiguity? It seems that we can get pretty accurate information about it once we get inside the graph.

torch/csrc/jit/passes/shape_analysis.cpp

ezyang

Nice work.

apaszke · 2018-08-24T15:30:35Z

We could track contiguity in TensorShape, but:

Contiguity is very tightly bound to strides, and not having to think about strides makes it much easier to work with the code. Most ops only have a size contract, but don't have a clear stride contract.
It doesn't seem to be necessary. I don't recall even a single place where it would enable us to do something we can't already do.

ezyang · 2018-08-24T15:36:24Z

Without symbolic sizes, we can't really track stride contracts in a useful way, but ops definitely have contiguity contracts which users know about, because contiguity tells you if you can view() a tensor, e.g.

But you're right, let's add it if/when a pass actually desperately wants to know about contiguity.

facebook-github-bot

apaszke has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

zdevito

This looks good -- I have a bunch of individual but minor comments and questions.

Higher level notes for building on top of this:

This puts the FusionCompiler at the breaking point of complexity. Further functionality added there is going require refactor some of the Chunk and Concat logic into a separate phase, given how many times we are need to check nSubtensors, and do other things to derive correct sizes. The core of the fusion is simple, but this chunk/concat stuff is getting spread all over the place.
We may notice regressions from overhead in launching fused kernels. For reference, we know the time it takes to do chunk in the interpreter adds significant overhead, and that seems on the same order of magnitude as the extra checking added here. We will need to monitor this and optimize if necessary.

torch/csrc/jit/autodiff.cpp

torch/csrc/jit/export.cpp

torch/csrc/jit/import.cpp

torch/csrc/jit/interpreter.cpp

torch/csrc/jit/passes/batch_mm.cpp

torch/csrc/jit/fusion_compiler.cpp

torch/csrc/jit/passes/shape_analysis.cpp

Most of the changes needed to be applied to the fuser, which heavily relied on this info. It now includes some extra run-time shape checks to determine if it can use its fused kernels, or if it should fall back to regular execution instead.

apaszke · 2018-08-25T18:50:11Z

The last commit forces slicing on types if the dynamic cast succeeds, but the kind doesn't match exactly. Normally we could simply use the copy constructor in this place, but because we're incorrectly comparing only addresses of types in many cases (because we assume they're used as singletons), the slicing needs some extra care. We really should either stop using shared pointers to types in most places (that would also come with a benefit of not having to incref/decref just to check the type), or have a subclass of shared pointer that uses equality on held elements in operator==.

facebook-github-bot

apaszke has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

* upstream/master: (89 commits) move HeatmapMaxKeypointOp unittest to oss fix xfails involving literals (pytorch#10905) Bag of Distributions doc fixes (pytorch#10894) Remove FIXME_zerol() from test_jit.py (pytorch#10900) Increase BC for PackedSequence ctor (pytorch#9864) Remove ability of Scalars to hold Tensors. Begin a bestiary of MSVC/NVCC bugs. (pytorch#10883) Prevent JIT from overspecializing to every single size configuration (pytorch#10844) Handling failing test on ROCm. Update mobile predictor caller's interface Cache isContiguous and numel Create class constant for string literal 'blob_names' Conv BN fusion for 3D conv (pytorch#10239) Stop using symbolic override for tracing RNNs (pytorch#10638) Add registry to pybind_state (pytorch#10759) Remove the nanopb submodule Create at::linear (pytorch#10799) Refactor THCNumerics and add common math functions for at::Half (pytorch#10301) Remove Tensor constructor of Scalar. (pytorch#10852) Revert D9492561: [pytorch][PR] Moving the operator argument to the front for kernelPointwiseApply. ...

…ytorch#10844) Summary: Please review the expects carefully to make sure there are no regressions. I tried to go over them one by one when they changed, but it's sometimes easy to miss finer details. Summary of changes: - Renamed `TensorType` to `CompleteTensorType`. Added a new `TensorType` which records only the scalar type, number of dimensions, and device of a value. The argument behind the rename is to encourage people to use `CompleteTensorType` less, as most passes will only have limited information available. To make transition easier `complete_type->cast<TensorType>()` works, and makes our passes work with both kinds of specialization if they don't need extra the extra detail. - Renamed `ArgumentSpec` to `CompleteArgumentSpec`. Added a new `ArgumentSpec`, which matches argument only at the level of the new `TensorType`. - Shape analysis can process graphs with both `CompleteTensorType` and `TensorType`. - Fuser was a part that heavily relied on full shape information being available. Now, we simply try to fuse the largest possible graphs, and have to do run-time checks to make sure they match the code we generate. If they don't, we fall back to regular interpretation. The shape checks are implementing using an optimized method exploiting algebraic properties of shapes with broadcasting, and the relations of broadcasting with pointwise ops. A full written proof of correctness of the shape checking algorithm is included in a comment in `graph_fuser.cpp`. zdevito ezyang mruberry ngimel csarofeen Pull Request resolved: pytorch#10844 Differential Revision: D9498705 Pulled By: apaszke fbshipit-source-id: 0c53c2fcebd871cc2a29c260f8d012276479cc61

apaszke requested review from Yangqing, anderspapitto, bddppq, colesbury, dzhulgakov, ezyang, gchanan, houseroad, jamesr66a, smessmer, soumith and zdevito as code owners August 24, 2018 01:45