mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
97eb7a281d
The big semantic change (and the reason for this port) is that we no longer monkeypatch Tensor with torchdim's special methods. The new algorithm for handling dispatch is that we first land in `__torch_function__` and we see if a special FCD implementation needs to be dispatch to first, and if there is nothing we fallback to the standard level strategy. Because there is no longer C binding equivalent of classes, we've condensed _C.Dim and Dim together, and similar for Tensor. This resulted in some bugs as the Python API is sometimes different from the C API. I've attempted to disambiguate these but there may still be mistakes (many early bugs were due to this problem). Dim and DimEntry are especially painful as Dim must abide by Tensor equality semantics, but is pointer equality in C (DimEntry doesn't have this problem). Another difference between C/Python that is subtle is we no longer get implicit conversions from Dim to DimEntry, this also caused some bugs. Much of the mechanical porting work was done by claude code. I have a separate PR that deletes functorch._C, but it was useful having dim.cpp to point claude at it so I haven't done it in this PR. From a reviewing perspective, I need to re-review that I didn't forget to port anything, some noticeably missing "small" things are patched_dim_method. I am still in progress of carefully doing a side-by-side review of ports; "simplifications" from claude code were also a major source of bugs. There are two major feature gaps in the implementation: - DelayedTensor and dot handling are not implemented yet. This should be reasonably easy, just need to do it. However, for the purposes of sharded propagation it is actually better not to reconstruct matmuls. - Splitting dimensions with an index like `[x, y]` doesn't work. The problem is that `__getitem__` interprets this as advanced indexing and sends the list to torch.tensor to turn into a tensor, instead of being eligible for `__torch_function__`. I think I might need to hard code a special case for this or something? Signed-off-by: Edward Yang <ezyang@meta.com> Pull Request resolved: https://github.com/pytorch/pytorch/pull/160236 Approved by: https://github.com/zdevito, https://github.com/albanD
48 lines
1.5 KiB
Python
48 lines
1.5 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the BSD-style license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
import os
|
|
import signal
|
|
import subprocess
|
|
from collections.abc import Generator
|
|
from contextlib import contextmanager
|
|
|
|
|
|
@contextmanager
|
|
def magic_trace(
|
|
output: str = "trace.fxt", magic_trace_cache: str = "/tmp/magic-trace"
|
|
) -> Generator[None, None, None]:
|
|
pid = os.getpid()
|
|
if not os.path.exists(magic_trace_cache):
|
|
print(f"Downloading magic_trace to: {magic_trace_cache}")
|
|
subprocess.run(
|
|
[
|
|
"wget",
|
|
"-O",
|
|
magic_trace_cache,
|
|
"-q",
|
|
"https://github.com/janestreet/magic-trace/releases/download/v1.0.2/magic-trace",
|
|
]
|
|
)
|
|
subprocess.run(["chmod", "+x", magic_trace_cache])
|
|
args = [magic_trace_cache, "attach", "-pid", str(pid), "-o", output]
|
|
p = subprocess.Popen(args, stderr=subprocess.PIPE, encoding="utf-8")
|
|
assert p.stderr is not None
|
|
while True:
|
|
x = p.stderr.readline()
|
|
print(x)
|
|
if "Attached" in x:
|
|
break
|
|
try:
|
|
yield
|
|
finally:
|
|
p.send_signal(signal.SIGINT)
|
|
r = p.wait()
|
|
if p.stderr is not None:
|
|
print(p.stderr.read())
|
|
p.stderr.close()
|
|
if r != 0:
|
|
raise ValueError(f"magic_trace exited abnormally: {r}")
|