SIGN IN SIGN UP

[PyTorch][jit] Avoid IValue copies in unpickler (#193250) (#193250)

Summary:

Avoid `IValue` copies in `unpickler.cpp`.
- Move keys/values into the dict in `readInstruction`'s `DICT`/`SETITEMS`/`SETITEM` arms. Those stack slots are dropped by `stack_.erase()` on the very next statement, and `Dict::insert_or_assign` is `template <class Key_, class Value_>(Key_&&, Value_&&)` forwarding into `Key(std::forward<Key_>(key))`, so this genuinely moves -- 4 atomic refcount ops saved per dict entry on the `state_dict` path.
- Change `const auto args_elems = args->elements()` to `const auto&` in `rebuildTensorFromTypeV2` and `rebuildParameter`. `auto` deduces `TupleElements` by value and deep-copies it: a heap allocation plus an `IValue` refcount bump per element once the tuple is longer than 3. `rebuildTensor` and `rebuildSparseTensor` in the same file already use `const auto&`.

No behavior change.

Test Plan:
Scratch `cpp_binary` against `//caffe2:torch-cpp`, built and run once at the parent commit and once with this diff. It covers `restoreAccurateTypeTags` over a class graph hitting every switch arm (generic list, dict-of-lists, tuple, present and absent `Optional`, nested object, tensor) plus aliased-handle, idempotence and shared-pointer cases; a `torch::jit::pickle`/`unpickle` round trip through the `DICT`, `SETITEMS` and `SETITEM` opcodes including an aliased value stored under two keys; and three Python-generated archives that drive `rebuildParameter` / `rebuildTensorFromTypeV2`.

Those last three exist because the C++ `Pickler` never emits `torch._utils._rebuild_parameter` or `torch._tensor._rebuild_from_type_v2`, so both functions are unreachable from a C++-only round trip. The fixtures are stdlib-`pickle` streams built against stub globals carrying the right `__module__`/`__qualname__`, spliced into a container emitted by `torch::pickle_save`. `param_multi` repeats `_rebuild_parameter`, so the repeats arrive as `BINGET` exactly as they do in a real archive, and its dict drives the `SETITEMS` path this diff moves out of.

```
$ diff /tmp/t_out_base.txt /tmp/t_out_change.txt
IDENTICAL

$ cat /tmp/t_out_change.txt
[restoreAccurateTypeTags] Object __torch__.E2E{lst=List<str>[a,b,c,],dct=Dict<str,List[int]>{k1:List<int>[1,2,],k2:List<int>[3,],},tup=Tuple(42,List<str>[t0,t1,],),opt=List<Any>[o0,],opt_none=None,nested=Object __torch__.Inner{inner_lst=List<str>[i0,i1,],},tensor=Tensor(sizes=[2, 3],strides=[3, 1],dtype=Float,sum=6,requires_grad=0,vals=[1,1,1,1,1,1,]),}
[restoreAccurateTypeTags x2] Object __torch__.E2E{lst=List<str>[a,b,c,],dct=Dict<str,List[int]>{k1:List<int>[1,2,],k2:List<int>[3,],},tup=Tuple(42,List<str>[t0,t1,],),opt=List<Any>[o0,],opt_none=None,nested=Object __torch__.Inner{inner_lst=List<str>[i0,i1,],},tensor=Tensor(sizes=[2, 3],strides=[3, 1],dtype=Float,sum=6,requires_grad=0,vals=[1,1,1,1,1,1,]),}
[aliased dict after traversal] Dict<str,List[int]>{k1:List<int>[1,2,],k2:List<int>[3,],}
[aliased inner after traversal] Object __torch__.Inner{inner_lst=List<str>[i0,i1,],}
[shared-pointer graph] Object __torch__.Shared{a=List<str>[s0,s1,],b=List<str>[s0,s1,],}
[pickle in ] Dict<str,Any>{inner:Dict<str,int>{key0:0,key10:70,key11:77,key1:7,key2:14,key3:21,key4:28,key5:35,key6:42,key7:49,key8:56,key9:63,},lst:List<Any>[1,two,],tensor:Tensor(sizes=[2, 3],strides=[3, 1],dtype=Long,sum=15,requires_grad=0,vals=[0,1,2,3,4,5,]),tup:Tuple(9,nine,),}
[pickle out] Dict<str,Any>{inner:Dict<str,int>{key0:0,key10:70,key11:77,key1:7,key2:14,key3:21,key4:28,key5:35,key6:42,key7:49,key8:56,key9:63,},lst:List<Any>[1,two,],tensor:Tensor(sizes=[2, 3],strides=[3, 1],dtype=Long,sum=15,requires_grad=0,vals=[0,1,2,3,4,5,]),tup:Tuple(9,nine,),}
[pickle bytes] 398
[pickle single] Dict<str,int>{solo:1,}
[pickle aliased] Dict<str,Any>{a:List<Any>[5,6,],b:List<Any>[5,6,],}
[fixture param.container] Tensor(sizes=[2, 3],strides=[3, 1],dtype=Float,sum=15,requires_grad=1,vals=[0,1,2,3,4,5,])
[fixture subclass.container] Tensor(sizes=[2, 3],strides=[3, 1],dtype=Float,sum=15,requires_grad=0,vals=[0,1,2,3,4,5,])
[fixture param_multi.container] Dict<Any,Any>{b:Tensor(sizes=[3],strides=[1],dtype=Float,sum=12,requires_grad=1,vals=[3,4,5,]),flat:Tensor(sizes=[6],strides=[1],dtype=Float,sum=15,requires_grad=0,vals=[0,1,2,3,4,5,]),w:Tensor(sizes=[2, 3],strides=[3, 1],dtype=Float,sum=15,requires_grad=1,vals=[0,1,2,3,4,5,]),}
[done]
```

`requires_grad=1` on `param.container` is `rebuildParameter` running -- it is the only thing that sets it. `subclass.container` reaching a correct tensor at all is `rebuildTensorFromTypeV2` running. `param_multi`s `b` is a view at storage offset 3, so `[3,4,5]` confirms the storage plumbing survived the move.

One line above is worth pre-empting: `opt=List<Any>[o0,]` sits next to `lst=List<str>[a,b,c,]`, so the list inside the `Optional` attribute does not get its tag refined. That is pre-existing and unrelated to this diff -- it reproduces identically on master -- and is now tracked as T284870639.

Build, with the TU proven to be in the compile path rather than a cache no-op:

```
$ buck2 build fbcode//mode/opt fbcode//caffe2:_libtorch
BUILD EXIT: 0

$ # same build with `#error` injected at the top of unpickler.cpp:
BUILD EXIT WITH #error: 3
Action failed: fbcode//caffe2:_libtorch (cxx_compile torch/csrc/jit/serialization/unpickler.cpp)
fbcode/caffe2/torch/csrc/jit/serialization/unpickler.cpp:1:2: error: SANITY_CHECK_TU_IS_COMPILED_V2
1 error generated.

$ # restored, rebuilt:
BUILD EXIT RESTORED: 0
```

`arc lint -e extra` on both copies reports nothing on any changed line. The findings it does report are all pre-existing, on lines this diff does not touch: `clang-diagnostic-shadow` at :257, `unreachable-code-break` at :665/:942/:982, `facebook-hte-DetailCall` at :1026, and `NULLSAFECLANG` on the `memcpy`s at :1161/:1171.

Not run: `fbcode//caffe2/test/cpp/jit:jit`. It is broken on this commit independently of this diff -- in `fbcode//mode/opt` the binary segfaults during `--gtest_list_tests` before any test body runs (exit 139, zero tests listed, reproduced with and without this change), and in `fbcode//mode/dev-nosan` it fails to configure with `fbcode//caffe2:mobile_bytecode is not visible to fbcode//caffe2/test/cpp/jit:jit-helper-lib`.

Reviewed By: dolpm

Differential Revision: D115809607

Pull Request resolved: https://github.com/pytorch/pytorch/pull/193250
Approved by: https://github.com/dolpm
J
Joshua Su committed
52d08b9519e82b71ababd98473077dff9007f875
Parent: dc3f3e4
Committed by PyTorch MergeBot <pytorchmergebot@users.noreply.github.com> on 8/25/2026, 1:11:03 AM