Respond to PR feedbackusers/RossBrunton/newMemInfo2

4 files changed, 18 insertions, 19 deletions

diff --git a/offload/include/omptarget.h b/offload/include/omptarget.h
index 197cbd3..4818ff6 100644
--- a/offload/include/omptarget.h
+++ b/offload/include/omptarget.h
@@ -98,7 +98,7 @@ enum OpenMPOffloadingDeclareTargetFlags {
 
 // Note: This type should be no larger than 3 bits, as the amdgpu platform uses
 // the lower 3 bits of a pointer to store it
-enum TargetAllocTy : int32_t {
+enum TargetAllocTy : uint8_t {
   TARGET_ALLOC_DEVICE = 0,
   TARGET_ALLOC_HOST,
   TARGET_ALLOC_SHARED,
diff --git a/offload/plugins-nextgen/amdgpu/src/rtl.cpp b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
index e06e4c6..7e6078d 100644
--- a/offload/plugins-nextgen/amdgpu/src/rtl.cpp
+++ b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
@@ -3575,15 +3575,11 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
 
     // The pointer info struct contains an "agent" field, but that doesn't
     // necessarily map to the device that created it
-    MemoryInfoTy ToReturn;
-    ToReturn.Base = Info.agentBaseAddress;
-    ToReturn.Size = Info.sizeInBytes;
     auto UserData = hsa_utils::UserDataPair::getFromOpaqueValue(Info.userData);
-    ToReturn.Type = static_cast<TargetAllocTy>(UserData.getInt());
-    ToReturn.Device =
-        reinterpret_cast<GenericDeviceTy *>(UserData.getPointer());
-
-    return ToReturn;
+    return MemoryInfoTy(
+        /*Base=*/Info.agentBaseAddress, /*Size=*/Info.sizeInBytes,
+        /*Type=*/static_cast<TargetAllocTy>(UserData.getInt()),
+        /*Device=*/reinterpret_cast<GenericDeviceTy *>(UserData.getPointer()));
   }
 
 private:
diff --git a/offload/plugins-nextgen/common/include/PluginInterface.h b/offload/plugins-nextgen/common/include/PluginInterface.h
index 163b02a..5f9d125 100644
--- a/offload/plugins-nextgen/common/include/PluginInterface.h
+++ b/offload/plugins-nextgen/common/include/PluginInterface.h
@@ -1277,6 +1277,10 @@ struct MemoryInfoTy {
   TargetAllocTy Type;
   GenericDeviceTy *Device;
 
+  MemoryInfoTy(void *Base, size_t Size, TargetAllocTy Type,
+               GenericDeviceTy *Device)
+      : Base(Base), Size(Size), Type(Type), Device(Device) {}
+
   void *limit() const { return reinterpret_cast<char *>(Base) + Size; }
 };
 
diff --git a/offload/plugins-nextgen/cuda/src/rtl.cpp b/offload/plugins-nextgen/cuda/src/rtl.cpp
index 2765453..f0343f5 100644
--- a/offload/plugins-nextgen/cuda/src/rtl.cpp
+++ b/offload/plugins-nextgen/cuda/src/rtl.cpp
@@ -1541,17 +1541,17 @@ struct CUDAPluginTy final : public GenericPluginTy {
 
     // Fast case, we have been given the base pointer directly
     if (MemoryInfo.contains(TgtPtr))
-      return MemoryInfo[TgtPtr];
+      return MemoryInfo.at(TgtPtr);
 
     // Slower case, we need to look up the base pointer first
     auto Loc = std::lower_bound(MemoryBases.begin(), MemoryBases.end(), TgtPtr,
                                 [&](const void *Iter, const void *Val) {
-                                  return MemoryInfo[Iter].limit() < Val;
+                                  return MemoryInfo.at(Iter).limit() < Val;
                                 });
-    if (Loc == MemoryBases.end() || TgtPtr > MemoryInfo[*Loc].limit())
+    if (Loc == MemoryBases.end() || TgtPtr > MemoryInfo.at(*Loc).limit())
       return Plugin::error(ErrorCode::NOT_FOUND,
                            "allocated memory information not found");
-    return MemoryInfo[*Loc];
+    return MemoryInfo.at(*Loc);
   }
 
 private:
@@ -1612,12 +1612,11 @@ void *CUDADeviceTy::allocate(size_t Size, void *, TargetAllocTy Kind) {
         std::lower_bound(CudaPlugin->MemoryBases.begin(),
                          CudaPlugin->MemoryBases.end(), MemAlloc),
         MemAlloc);
-    CudaPlugin->MemoryInfo[MemAlloc] = MemoryInfoTy{
-        /*Base=*/MemAlloc,
-        /*Size=*/Size,
-        /*Type=*/Kind,
-        /*Device=*/this,
-    };
+    CudaPlugin->MemoryInfo.emplace_or_assign(MemAlloc,
+                                             /*Base=*/MemAlloc,
+                                             /*Size=*/Size,
+                                             /*Type=*/Kind,
+                                             /*Device=*/this);
   }
   return MemAlloc;
 }