pytorch
e936277c - [ROCm] force HIP context initialization for inductor UTs (#103149)

Commit
1 year ago
[ROCm] force HIP context initialization for inductor UTs (#103149) Workaround for https://github.com/pytorch/pytorch/issues/102886 related to: https://github.com/pytorch/pytorch/issues/102476 https://github.com/pytorch/pytorch/issues/102475 https://github.com/pytorch/pytorch/issues/102474 https://github.com/pytorch/pytorch/issues/102473 https://github.com/pytorch/pytorch/issues/102473 https://github.com/pytorch/pytorch/issues/102472 Since https://github.com/pytorch/pytorch/commit/9aaa12e32855a711ba6e3135138a46e63402cbbe the first inductor (CPU) UT fails until the GPU context is correct initialised and the subsequent UTs pass. CUDA observes the same issue and a workaround was pushed to force initialisation of cuda context by declaring an empty tensor https://github.com/pytorch/pytorch/issues/92627, we have adopted the same approach but have opted for `torch.zeros` which correctly activates the HIP context after the kernel launch. **Reproducer:** ``` import torch from torch._subclasses.fake_tensor import FakeTensorMode import argparse if __name__ == '__main__': parser = argparse.ArgumentParser(description='Swap between torch.empty and torch.randn operations.') parser.add_argument('--empty', action='store_true', help='Use torch.empty operation') parser.add_argument('--rand', action='store_true', help='Use torch.randn operation') args = parser.parse_args() torch.cuda.set_device(0) if args.empty: torch.empty(1, device="cuda") elif args.rand: torch.rand(1, device="cuda") print(f": hasPrimaryContext: {torch._C._cuda_hasPrimaryContext(0)") with FakeTensorMode(): p = torch.randn(4, 2, requires_grad=True, device='cuda') x = torch.randn(8, 4, device='cuda') y = torch.mm(x, p).square().sum() y.backward() ``` **ROCm python repro.py --empty** 0: hasPrimaryContext: False **ROCm python repro.py --rand** 0: hasPrimaryContext: True **CUDA python repro.py --empty** 0: hasPrimaryContext: True **CUDA python repro.py --rand** 0: hasPrimaryContext: True Pull Request resolved: https://github.com/pytorch/pytorch/pull/103149 Approved by: https://github.com/eellison
Author
Committer
Parents
Loading