frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-21 05:34:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Inductor] Inplacing with Donated Buffer (#140113)

Browse Source 
Currently, inductor does not inplace update a buffer if it is an input buffer. Because we don't know if an input will be used by other functions.

Donated buffer provides additional information that an input buffer will not be used by other functions. So we can inplace update donated buffer when possible.

[Dashboard](https://hud.pytorch.org/benchmark/torchbench/inductor_dynamic?dashboard=torchinductor&startTime=Mon,%2011%20Nov%202024%2018:14:36%20GMT&stopTime=Mon,%2018%20Nov%202024%2018:14:36%20GMT&granularity=hour&mode=training&dtype=amp&deviceName=cuda%20(a100)&lBranch=bf/donated-buffer-inplace&lCommit=5df0769c00e6f9000caeb10fd5cbf0b165f69c2a&rBranch=main&rCommit=2b39a8db7741b816b03677a9c6fec1af05640dee)

![image](https://github.com/user-attachments/assets/f19d961f-7973-418e-9de8-5c2a97950478)
![image](https://github.com/user-attachments/assets/df3bd6a9-58b8-4e8a-8397-9e3b1de9adfe)

Pull Request resolved: https://github.com/pytorch/pytorch/pull/140113
Approved by: https://github.com/eellison
This commit is contained in:
Boyuan Feng
2024-11-26 17:19:50 +00:00
committed by PyTorch MergeBot
parent 3ef031909f
commit eecc8e362c
7 changed files with 180 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
torch/_inductor/codegen/wrapper.py
View File

@ -2120,7 +2120,11 @@ class PythonWrapperCodegen(CodeGen):
def codegen_allocation(self, buffer: ir.Buffer):
name = buffer.get_name()
if name in V.graph.removed_buffers or name in self.allocated:
if (
name in V.graph.removed_buffers
or name in self.allocated
or isinstance(buffer, ir.DonatedBuffer)
):
return
self.allocated.add(name)
if isinstance(
@ -2174,7 +2178,12 @@ class PythonWrapperCodegen(CodeGen):
name = input_buffer.get_name()
return not (
name in V.graph.removed_buffers
or name in V.graph.graph_inputs
or (
name in V.graph.graph_inputs
and not isinstance(
V.graph.graph_inputs_original[name], ir.DonatedBuffer
)
)
or name in V.graph.constants
or name in V.graph.torchbind_constants
or name in V.graph.never_reuse_buffers