[Offload] Allow querying the size of globals

The `GlobalTy` helper has been extended to make both the Size and Ptr be
optional. Now `getGlobalMetadataFromDevice`/`Image` is able to write the
size of the global to the struct, instead of just verifying it.

Author: RossBrunton

offload/plugins-nextgen/amdgpu/src/rtl.cpp

@@ -3089,15 +3089,16 @@ struct AMDGPUGlobalHandlerTy final : public GenericGlobalHandlerTy {
     }
 
     // Check the size of the symbol.
-    if (SymbolSize != DeviceGlobal.getSize())
+    if (DeviceGlobal.hasSize() && SymbolSize != DeviceGlobal.getSize())
       return Plugin::error(
           ErrorCode::INVALID_BINARY,
           "failed to load global '%s' due to size mismatch (%zu != %zu)",
           DeviceGlobal.getName().data(), SymbolSize,
           (size_t)DeviceGlobal.getSize());
 
-    // Store the symbol address on the device global metadata.
+    // Store the symbol address and size on the device global metadata.
     DeviceGlobal.setPtr(reinterpret_cast<void *>(SymbolAddr));
+    DeviceGlobal.setSize(SymbolSize);
 
     return Plugin::success();
   }

offload/plugins-nextgen/common/include/GlobalHandler.h

@@ -37,20 +37,33 @@ using namespace llvm::object;
 /// Common abstraction for globals that live on the host and device.
 /// It simply encapsulates the symbol name, symbol size, and symbol address
 /// (which might be host or device depending on the context).
+/// Both size and address may be absent, and can be populated with
+// getGlobalMetadataFromDevice/Image.
 class GlobalTy {
   // NOTE: Maybe we can have a pointer to the offload entry name instead of
   // holding a private copy of the name as a std::string.
   std::string Name;
-  uint32_t Size;
-  void *Ptr;
+  std::optional<uint32_t> Size;
+  std::optional<void *> Ptr;
 
 public:
-  GlobalTy(const std::string &Name, uint32_t Size, void *Ptr = nullptr)
+  GlobalTy(const std::string &Name) : Name(Name) {}
+  GlobalTy(const std::string &Name, uint32_t Size) : Name(Name), Size(Size) {}
+  GlobalTy(const std::string &Name, uint32_t Size, void *Ptr)
       : Name(Name), Size(Size), Ptr(Ptr) {}
 
   const std::string &getName() const { return Name; }
-  uint32_t getSize() const { return Size; }
-  void *getPtr() const { return Ptr; }
+  uint32_t getSize() const {
+    assert(hasSize() && "Size not initialised");
+    return *Size;
+  }
+  void *getPtr() const {
+    assert(hasPtr() && "Ptr not initialised");
+    return *Ptr;
+  }
+
+  bool hasSize() const { return Size.has_value(); }
+  bool hasPtr() const { return Ptr.has_value(); }
 
   void setSize(int32_t S) { Size = S; }
   void setPtr(void *P) { Ptr = P; }
@@ -139,8 +152,11 @@ class GenericGlobalHandlerTy {
   bool isSymbolInImage(GenericDeviceTy &Device, DeviceImageTy &Image,
                        StringRef SymName);
 
-  /// Get the address and size of a global in the image. Address and size are
-  /// return in \p ImageGlobal, the global name is passed in \p ImageGlobal.
+  /// Get the address and size of a global in the image. Address is
+  /// returned in \p ImageGlobal and the global name is passed in \p
+  /// ImageGlobal. If no size is present in \p ImageGlobal, then the size of the
+  /// global will be stored there. If it is present, it will be validated
+  /// against the real size of the global.
   Error getGlobalMetadataFromImage(GenericDeviceTy &Device,
                                    DeviceImageTy &Image, GlobalTy &ImageGlobal);
 
@@ -149,9 +165,11 @@ class GenericGlobalHandlerTy {
   Error readGlobalFromImage(GenericDeviceTy &Device, DeviceImageTy &Image,
                             const GlobalTy &HostGlobal);
 
-  /// Get the address and size of a global from the device. Address is return in
-  /// \p DeviceGlobal, the global name and expected size are passed in
-  /// \p DeviceGlobal.
+  /// Get the address and size of a global from the device. Address is
+  /// returned in \p ImageGlobal and the global name is passed in \p
+  /// ImageGlobal. If no size is present in \p ImageGlobal, then the size of the
+  /// global will be stored there. If it is present, it will be validated
+  /// against the real size of the global.
   virtual Error getGlobalMetadataFromDevice(GenericDeviceTy &Device,
                                             DeviceImageTy &Image,
                                             GlobalTy &DeviceGlobal) = 0;

offload/plugins-nextgen/cuda/src/rtl.cpp

@@ -1355,13 +1355,15 @@ class CUDAGlobalHandlerTy final : public GenericGlobalHandlerTy {
                                  GlobalName))
       return Err;
 
-    if (CUSize != DeviceGlobal.getSize())
+    if (DeviceGlobal.hasSize() && CUSize != DeviceGlobal.getSize())
       return Plugin::error(
           ErrorCode::INVALID_BINARY,
           "failed to load global '%s' due to size mismatch (%zu != %zu)",
           GlobalName, CUSize, (size_t)DeviceGlobal.getSize());
 
     DeviceGlobal.setPtr(reinterpret_cast<void *>(CUPtr));
+    DeviceGlobal.setSize(CUSize);
+
     return Plugin::success();
   }
 };