Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-18 16:54:33 +02:00
Code Issues Packages Projects Releases Wiki Activity

zluda hijack fft

Browse Source
This commit is contained in:
Seunghoon Lee
2024-09-24 11:01:48 +09:00
parent 4de0480d9f
commit 718cfe48a2
1 changed files with 15 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/zluda_hijacks.py
+15 -3
View File
@@ -3,12 +3,22 @@ from modules import rocm
_topk = torch.topk
def topk(tensor: torch.Tensor, *args, **kwargs):
device = tensor.device
values, indices = _topk(tensor.cpu(), *args, **kwargs)
def topk(input: torch.Tensor, *args, **kwargs): # pylint: disable=redefined-builtin
device = input.device
values, indices = _topk(input.cpu(), *args, **kwargs)
return torch.return_types.topk((values.to(device), indices.to(device),))
_fft_fftn = torch.fft.fftn
def fft_fftn(input: torch.Tensor, *args, **kwargs) -> torch.Tensor: # pylint: disable=redefined-builtin
return _fft_fftn(input.cpu(), *args, **kwargs).to(input.device)
_fft_ifftn = torch.fft.ifftn
def fft_ifftn(input: torch.Tensor, *args, **kwargs) -> torch.Tensor: # pylint: disable=redefined-builtin
return _fft_ifftn(input.cpu(), *args, **kwargs).to(input.device)
def jit_script(f, *_, **__): # experiment / provide dummy graph
f.graph = torch._C.Graph() # pylint: disable=protected-access
return f
@@ -17,4 +27,6 @@ def jit_script(f, *_, **__): # experiment / provide dummy graph
def do_hijack():
torch.version.hip = rocm.version
torch.topk = topk
torch.fft.fftn = fft_fftn
torch.fft.ifftn = fft_ifftn
torch.jit.script = jit_script