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
68 lines
2.1 KiB
Python
68 lines
2.1 KiB
Python
import dis
|
|
from typing import Any, Optional
|
|
|
|
|
|
class _PyInstDecoder:
|
|
"""
|
|
Decodes Python bytecode instructions to extract variable names
|
|
"""
|
|
|
|
def __init__(self, code_object: Any, lasti: int) -> None:
|
|
self.code_object = code_object
|
|
self.instructions = list(dis.get_instructions(code_object))
|
|
self.offset = self._find_instruction_index(lasti)
|
|
|
|
def _find_instruction_index(self, lasti: int) -> int:
|
|
"""Find instruction index corresponding to lasti (byte offset)."""
|
|
# Find the instruction at or before lasti
|
|
# This should find the CALL instruction, not the next one
|
|
best_idx = 0
|
|
for i, instr in enumerate(self.instructions):
|
|
if instr.offset <= lasti:
|
|
best_idx = i
|
|
else:
|
|
break
|
|
return best_idx
|
|
|
|
def next(self) -> None:
|
|
"""Advance to the next instruction."""
|
|
self.offset += 1
|
|
|
|
def opcode(self) -> Optional[str]:
|
|
"""Get the opcode name of the current instruction."""
|
|
if self.offset < len(self.instructions):
|
|
return self.instructions[self.offset].opname
|
|
return None
|
|
|
|
def oparg(self) -> int:
|
|
"""Get the argument of the current instruction."""
|
|
if self.offset < len(self.instructions):
|
|
return self.instructions[self.offset].arg or 0
|
|
return 0
|
|
|
|
def name(self) -> Optional[str]:
|
|
"""
|
|
Extract variable name from current instruction.
|
|
"""
|
|
opname = self.opcode()
|
|
if not opname:
|
|
return None
|
|
|
|
names = None
|
|
if opname in ("STORE_NAME", "STORE_GLOBAL"):
|
|
names = self.code_object.co_names
|
|
elif opname == "STORE_FAST":
|
|
names = self.code_object.co_varnames
|
|
elif opname == "STORE_DEREF":
|
|
names = self.code_object.co_cellvars
|
|
if not names:
|
|
names = self.code_object.co_freevars
|
|
else:
|
|
return None
|
|
|
|
arg = self.oparg()
|
|
if names and 0 <= arg < len(names):
|
|
return names[arg]
|
|
|
|
return None
|