[llvm] [OFFLOAD] Add interface to extend image validation (PR #185663)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-offload

Author: Alex Duran (adurang)

<details>
<summary>Changes</summary>

As discussed in #<!-- -->185404 we might want to provide a way for plugins to validate images not recognized by the common layer.

This PR adds such extension and uses to validate pure SPIRV images by the Level Zero plugin.

---
Full diff: https://github.com/llvm/llvm-project/pull/185663.diff


4 Files Affected:

- (modified) offload/plugins-nextgen/common/include/PluginInterface.h (+18) 
- (modified) offload/plugins-nextgen/common/src/PluginInterface.cpp (+8-2) 
- (modified) offload/plugins-nextgen/level_zero/include/L0Plugin.h (+2) 
- (modified) offload/plugins-nextgen/level_zero/src/L0Plugin.cpp (+7) 


``````````diff
diff --git a/offload/plugins-nextgen/common/include/PluginInterface.h b/offload/plugins-nextgen/common/include/PluginInterface.h
index 1c59ed1eda8416..6da8389636077f 100644
--- a/offload/plugins-nextgen/common/include/PluginInterface.h
+++ b/offload/plugins-nextgen/common/include/PluginInterface.h
@@ -1364,6 +1364,24 @@ struct GenericPluginTy {
   virtual Expected<bool> isELFCompatible(uint32_t DeviceID,
                                          StringRef Image) const = 0;
 
+  /// Indicate if an image is compatible with the plugin. This is called if
+  /// the image is not recognized as compatible by the common layer. This gives
+  /// the plugin a chance to inspect the image and decide if it is compatible.
+  virtual Expected<bool> isCompatibleImage(StringRef Image) const {
+    return false;
+  }
+
+  /// Indicate if an imaget is compatible with the plugin devices. This is
+  /// called if the image is not recognized as compatible by the common layer.
+  /// This gives the plugin a chance to inspect the image and decide if it is
+  /// compatible. Notice that this function may be called before actually
+  /// initializing the devices. So we could not move this function into
+  /// GenericDeviceTy.
+  virtual Expected<bool> isCompatibleImage(uint32_t DeviceID,
+                                           StringRef Image) const {
+    return isCompatibleImage(Image);
+  }
+
   virtual Error flushQueueImpl(omp_interop_val_t *Interop) {
     return Plugin::success();
   }
diff --git a/offload/plugins-nextgen/common/src/PluginInterface.cpp b/offload/plugins-nextgen/common/src/PluginInterface.cpp
index f681213b387945..0cc1c38e5481a3 100644
--- a/offload/plugins-nextgen/common/src/PluginInterface.cpp
+++ b/offload/plugins-nextgen/common/src/PluginInterface.cpp
@@ -1652,7 +1652,10 @@ int32_t GenericPluginTy::isPluginCompatible(StringRef Image) {
     return *MatchOrErr;
   }
   default:
-    return false;
+    auto MatchOrErr = isCompatibleImage(Image);
+    if (Error Err = MatchOrErr.takeError())
+      return HandleError(std::move(Err));
+    return *MatchOrErr;
   }
 }
 
@@ -1689,7 +1692,10 @@ int32_t GenericPluginTy::isDeviceCompatible(int32_t DeviceId, StringRef Image) {
     return *MatchOrErr;
   }
   default:
-    return false;
+    auto MatchOrErr = isCompatibleImage(DeviceId, Image);
+    if (Error Err = MatchOrErr.takeError())
+      return HandleError(std::move(Err));
+    return *MatchOrErr;
   }
 }
 
diff --git a/offload/plugins-nextgen/level_zero/include/L0Plugin.h b/offload/plugins-nextgen/level_zero/include/L0Plugin.h
index cd964a0d468907..709d2d8d4e26e9 100644
--- a/offload/plugins-nextgen/level_zero/include/L0Plugin.h
+++ b/offload/plugins-nextgen/level_zero/include/L0Plugin.h
@@ -107,6 +107,8 @@ class LevelZeroPluginTy final : public GenericPluginTy {
   Error flushQueueImpl(omp_interop_val_t *Interop) override;
   Error syncBarrierImpl(omp_interop_val_t *Interop) override;
   Error asyncBarrierImpl(omp_interop_val_t *Interop) override;
+
+  Expected<bool> isCompatibleImage(StringRef Image) const override;
 };
 
 } // namespace llvm::omp::target::plugin
diff --git a/offload/plugins-nextgen/level_zero/src/L0Plugin.cpp b/offload/plugins-nextgen/level_zero/src/L0Plugin.cpp
index 285fe797b5d7df..9a70ca3cfdf619 100644
--- a/offload/plugins-nextgen/level_zero/src/L0Plugin.cpp
+++ b/offload/plugins-nextgen/level_zero/src/L0Plugin.cpp
@@ -233,6 +233,13 @@ Error LevelZeroPluginTy::asyncBarrierImpl(omp_interop_val_t *Interop) {
   return Plugin::success();
 }
 
+// We only need to check for formats other than ELF here
+Expected<bool> LevelZeroPluginTy::isCompatibleImage(StringRef Image) const {
+  if (identify_magic(Image) == file_magic::spirv_object)
+    return true;
+  return false; 
+}
+
 } // namespace llvm::omp::target::plugin
 
 extern "C" {

``````````

</details>


https://github.com/llvm/llvm-project/pull/185663