From dce017da50c6f6fb49ca6b27c85dc8bea3ff92a0 Mon Sep 17 00:00:00 2001
From: Arash Pakbin <arash.pakbin@amd.com>
Date: Thu, 29 May 2025 21:12:09 +0000
Subject: [PATCH] [ROCm] MIOpen: Get current device from Torch rather than HIP
 in handle creation (#154549)

Get current device from Torch rather than HIP in MIOpen handle creation. The device may have already been set from torch side, otherwise device is set to 0 for handle.  Additional audits of cudnn vs miopen Handle.cpp file.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/154549
Approved by: https://github.com/jeffdaily, https://github.com/cyyever

Co-authored-by: Jeff Daily <jeff.daily@amd.com>
---
 aten/src/ATen/miopen/Handle.cpp | 51 ++++++++++++++++++---------------
 1 file changed, 28 insertions(+), 23 deletions(-)

diff --git a/aten/src/ATen/miopen/Handle.cpp b/aten/src/ATen/miopen/Handle.cpp
index 6b8c7c6421c4..7036c2e46a56 100644
--- a/aten/src/ATen/miopen/Handle.cpp
+++ b/aten/src/ATen/miopen/Handle.cpp
@@ -1,9 +1,11 @@
-#include <ATen/miopen/Exceptions.h>
-#include <ATen/miopen/Handle.h>
 #include <ATen/hip/detail/DeviceThreadHandles.h>
+#include <ATen/miopen/Handle.h>
 #include <c10/hip/HIPStream.h>
 
-namespace at { namespace native {
+#include <ATen/hip/Exceptions.h>
+#include <ATen/miopen/Exceptions.h>
+
+namespace at::native {
 namespace {
 
 void createMIOpenHandle(miopenHandle_t *handle) {
@@ -11,30 +13,33 @@ void createMIOpenHandle(miopenHandle_t *handle) {
 }
 
 void destroyMIOpenHandle(miopenHandle_t handle) {
-// this is because of something dumb in the ordering of
-// destruction. Sometimes atexit, the cuda context (or something)
-// would already be destroyed by the time this gets destroyed. It
-// happens in fbcode setting. @colesbury and I decided to not destroy
-// the handle as a workaround.
-//   - @soumith
-//
-// Further note: this is now disabled globally, because we are seeing
-// the same issue as mentioned above in CUDA 11 CI.
-//   - @zasdfgbnm
-//
-// #ifdef NO_MIOPEN_DESTROY_HANDLE
-// #else
-//   miopenDestroy(handle);
-// #endif
+  // this is because of something dumb in the ordering of
+  // destruction. Sometimes atexit, the cuda context (or something)
+  // would already be destroyed by the time this gets destroyed. It
+  // happens in fbcode setting. @colesbury and I decided to not destroy
+  // the handle as a workaround.
+  //   - @soumith
+  //
+  // Further note: this is now disabled globally, because we are seeing
+  // the same issue as mentioned above in CUDA 11 CI.
+  //   - @zasdfgbnm
+  //
+  // #ifdef NO_MIOPEN_DESTROY_HANDLE
+  // #else
+  //   miopenDestroy(handle);
+  // #endif
 }
 
-using MIOpenPoolType = at::cuda::DeviceThreadHandlePool<miopenHandle_t, createMIOpenHandle, destroyMIOpenHandle>;
+using MIOpenPoolType = at::cuda::DeviceThreadHandlePool<
+    miopenHandle_t,
+    createMIOpenHandle,
+    destroyMIOpenHandle>;
 
 } // namespace
 
 miopenHandle_t getMiopenHandle() {
-  int device;
-  HIP_CHECK(hipGetDevice(&device));
+  c10::DeviceIndex device = 0;
+  AT_CUDA_CHECK(c10::hip::GetDevice(&device));
 
   // Thread local PoolWindows are lazily-initialized
   // to avoid initialization issues that caused hangs on Windows.
@@ -46,8 +51,8 @@ miopenHandle_t getMiopenHandle() {
       pool->newPoolWindow());
 
   auto handle = myPoolWindow->reserve(device);
-  MIOPEN_CHECK(miopenSetStream(handle, at::hip::getCurrentHIPStream()));
+  MIOPEN_CHECK(miopenSetStream(handle, c10::hip::getCurrentHIPStream()));
   return handle;
 }
 
-}} // namespace at::native
+} // namespace at::native