[mlir][ExecutionEngine] Remove ScopedContext from ROCm wrappers

[krzysz00]

Summary

The push/pop context APIs are deprecated in HIP, and keeping the
default device set is handled in IHP using hipSetDevice().

mlir/lib/ExecutionEngine/RocmRuntimeWrappers.cpp

[](hipError_t result) {                                                      \
     const char *name = hipGetErrorName(result);                                \
     if (!name)                                                                 \
       name = "<unknown>";                                                      \
     fprintf(stderr, "'%s' failed with '%s'\n", #expr, name);                   \
   }(expr)
 
 thread_local static int32_t defaultDevice = 0;
 
-// Sets the `Context` for the duration of the instance and restores the previous
-// context on destruction.
-class ScopedContext {
-public:
-  ScopedContext() {
-    // Static reference to HIP primary context for device ordinal defaultDevice.
-    static hipCtx_t context = [] {
-      HIP_REPORT_IF_ERROR(hipInit(/*flags=*/0));
-      hipDevice_t device;
-      HIP_REPORT_IF_ERROR(hipDeviceGet(&device, /*ordinal=*/defaultDevice));
-      hipCtx_t ctx;
-      HIP_REPORT_IF_ERROR(hipDevicePrimaryCtxRetain(&ctx, device));
-      return ctx;
-    }();
-
-    HIP_REPORT_IF_ERROR(hipCtxPushCurrent(context));
-  }
-
-  ~ScopedContext() { HIP_REPORT_IF_ERROR(hipCtxPopCurrent(nullptr)); }
-};
-
 extern "C" hipModule_t mgpuModuleLoad(void *data) {
-  ScopedContext scopedContext;
   hipModule_t module = nullptr;
   HIP_REPORT_IF_ERROR(hipModuleLoadData(&module, data));
   return module;
 }
 
 extern "C" void mgpuModuleUnload(hipModule_t module) {
   HIP_REPORT_IF_ERROR(hipModuleUnload(module));
 }

 // the type of MLIR's index type. This avoids the need for casts in the
 // generated MLIR code.
 extern "C" void mgpuLaunchKernel(hipFunction_t function, intptr_t gridX,
                                  intptr_t gridY, intptr_t gridZ,
                                  intptr_t blockX, intptr_t blockY,
                                  intptr_t blockZ, int32_t smem,
                                  hipStream_t stream, void **params,
                                  void **extra) {
-  ScopedContext scopedContext;
   HIP_REPORT_IF_ERROR(hipModuleLaunchKernel(function, gridX, gridY, gridZ,
                                             blockX, blockY, blockZ, smem,
                                             stream, params, extra));
 }
 
 extern "C" hipStream_t mgpuStreamCreate() {
-  ScopedContext scopedContext;
   hipStream_t stream = nullptr;
   HIP_REPORT_IF_ERROR(hipStreamCreate(&stream));
   return stream;
 }
 
 extern "C" void mgpuStreamDestroy(hipStream_t stream) {
   HIP_REPORT_IF_ERROR(hipStreamDestroy(stream));
 }
 
 extern "C" void mgpuStreamSynchronize(hipStream_t stream) {
   return HIP_REPORT_IF_ERROR(hipStreamSynchronize(stream));
 }
 
 extern "C" void mgpuStreamWaitEvent(hipStream_t stream, hipEvent_t event) {
   HIP_REPORT_IF_ERROR(hipStreamWaitEvent(stream, event, /*flags=*/0));
 }
 
 extern "C" hipEvent_t mgpuEventCreate() {
-  ScopedContext scopedContext;
   hipEvent_t event = nullptr;
   HIP_REPORT_IF_ERROR(hipEventCreateWithFlags(&event, hipEventDisableTiming));
   return event;
 }
 
 extern "C" void mgpuEventDestroy(hipEvent_t event) {
   HIP_REPORT_IF_ERROR(hipEventDestroy(event));
 }
 
 extern "C" void mgpuEventSynchronize(hipEvent_t event) {
   HIP_REPORT_IF_ERROR(hipEventSynchronize(event));
 }
 
 extern "C" void mgpuEventRecord(hipEvent_t event, hipStream_t stream) {
   HIP_REPORT_IF_ERROR(hipEventRecord(event, stream));
 }
 
 extern "C" void *mgpuMemAlloc(uint64_t sizeBytes, hipStream_t /*stream*/) {
-  ScopedContext scopedContext;
   void *ptr;
   HIP_REPORT_IF_ERROR(hipMalloc(&ptr, sizeBytes));
   return ptr;
 }
 
 extern "C" void mgpuMemFree(void *ptr, hipStream_t /*stream*/) {
   HIP_REPORT_IF_ERROR(hipFree(ptr));
 }
extern "C" void mgpuMemset32(void *dst, int value, size_t count,
   HIP_REPORT_IF_ERROR(hipMemsetD32Async(reinterpret_cast<hipDeviceptr_t>(dst),
                                         value, count, stream));
 }
 /// Helper functions for writing mlir example code
 
 // Allows to register byte array with the ROCM runtime. Helpful until we have
 // transfer functions implemented.
 extern "C" void mgpuMemHostRegister(void *ptr, uint64_t sizeBytes) {
-  ScopedContext scopedContext;
   HIP_REPORT_IF_ERROR(hipHostRegister(ptr, sizeBytes, /*flags=*/0));
 }
 
 // Allows to register a MemRef with the ROCm runtime. Helpful until we have
 // transfer functions implemented.
 extern "C" void
 mgpuMemHostRegisterMemRef(int64_t rank, StridedMemRefType<char, 1> *descriptor,
                           int64_t elementSizeBytes) {