[OpenMP][ClangLinkerWrapper] Extending linker wrapper to embed metadata for multi-arch fat binaries

[saiislam]

Summary

This patch adds "__tgt_image_info" field for each of the images
embedded in a multi-arch image. Required changes in libomptarget
are also shown.

The information in "tgt_image_info" struct is provided in the clang-linker-wrapper
as a call to tgt_register_image_info for each image in the library
of images also created by the clang-linker-wrapper.
tgt_register_image_info is called for each image BEFORE the single
call to tgt_register_lib so that image information is available
before they are loaded. clang-linker-wrapper gets this image information
from command line arguments provided by the clang driver when it creates
the call to the clang-linker-wrapper command.
This architecture allows the binary image (pointed to by ImageStart and
ImageEnd in tgt_device_image) to remain architecture indenendent.
That is, the architecture independent part of the libomptarget runtime
does not need to peer inside the image to determine if it is loadable
even though in most cases the image is an elf object.
There is one tgt_image_info for each tgt_device_image. For backward
compabibility, no changes are allowed to either tgt_device_image or
tgt_bin_desc. The absense of __tgt_image_info is the indication that
the runtime is being used on a binary created by an old version of
the compiler.

[jdoerfert]

openmp/libomptarget/include/omptarget.h

197

I doubt we want/need to number and count them. This seems fragile and unhelpful. "initial allocation" is not expensive, image numbers should be implicit.

[jhuber6]

clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp

1100–1103

I'm doing something similar in D123810, I just used the existing DeviceFile because I needed the Arch and Kind fields to dispatch the appropriate wrapping job for CUDA / HIP.

clang/tools/clang-linker-wrapper/OffloadWrapper.cpp

285

I'm wondering if it would be better to create a new __tgt_bin_desc and call a new __tgt_register_lib with it here so we don't need multiple calls here. Inside that new runtime function we could just widen or shrink the existing structs as needed. That way each device image would have this metadata associated with it and the target plugin can handle it as-needed.

[yaxunl]

need a test for the generated registration code

clang/tools/clang-linker-wrapper/OffloadWrapper.cpp

120–126

I am wondering whether we should add a few more fields to make it more generic for all offloading languages and platforms:

char* target_triple;
char* offloading_kind; // openmp, hip, etc
char* file_type; // elf, spirv, bitcode, etc

[saiislam]

In D124525#3478469, @yaxunl wrote:

need a test for the generated registration code

Yes, I will add tests.

clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp

1100–1103

Seems simpler. I will pull that change here.

clang/tools/clang-linker-wrapper/OffloadWrapper.cpp

120–126

Good idea. Though I am not sure whether file_type info is being propagated in by the linker-wrapper or not. I will check.

285

Last time multiple vendors objected to changing __tgt_bin_desc and __tgt_device_image structs. The reason was backward compatibility of multiple downstream runtimes.

[jhuber6]

clang/tools/clang-linker-wrapper/OffloadWrapper.cpp

285

If you keep the old __tgt_register_lib function, old applications will use that. Then inside libomptarget you can change __tgt_register_lib function to create a new __tgt_device_image from the old format. It would have some time penalty as we'd need to allocate some new structs, but it might be easier to deal with if each image contained its own metadata for the plugin to use. I was imagining it would look something like

struct __tgt_device_image_v2 {
  void* 	ImageStart;
  void* 	ImageEnd;
  __tgt_offload_entry* 	EntriesBegin;
  __tgt_offload_entry* 	EntriesEnd;
  __tgt_image_into *  ImageInfo;
};

Not sure if this is the best approach, but it should let us keep backwards compatibility and change the structs.

[yaxunl]

clang/tools/clang-linker-wrapper/OffloadWrapper.cpp

120–126

the file_type is for the file embedded by OffloadWrapper, therefore OffloadWrapper knows it.

Currently it is fatbin or cubin for CUDA and code object (elf) for HIP.

In the future, it may be spirv or bitcode to allow JIT compilation of spirv or bitcode in runtime so that one binary for all GPU's. OffloadWrapper will decide what is the final file type to embed based on type of embedded binaries in input files, target triple, and compile arguments.

[yaxunl]

clang/tools/clang-linker-wrapper/OffloadWrapper.cpp

120–126

the file_type is for the file embedded by OffloadWrapper, therefore OffloadWrapper knows it.

Currently it is fatbin or cubin for CUDA and code object (elf) for HIP.

In the future, it may be spirv or bitcode to allow JIT compilation of spirv or bitcode in runtime so that one binary for all GPU's. OffloadWrapper will decide what is the final file type to embed based on type of embedded binaries in input files, target triple, and compile arguments.

Sorry I mean ClangLinkerWrapper. ClangLinkerWrapper should somehow pass the file type info to OffloadWrapper.

[jhuber6]

I'm suggesting instead we define a new __tgt_device_image and a new __tgt_register_lib function to support this. This new __tgt_device_image will simply contain a pointer to an optional information struct.

struct __tgt_device_image_v2 {
  void* 	ImageStart;
  void* 	ImageEnd;
  __tgt_offload_entry* 	EntriesBegin;
  __tgt_offload_entry* 	EntriesEnd;
  __tgt_image_into*  ImageInfo;
};

This new struct breaks the ABI with the old __tgt_device_image because these are put into an array and we change the size, but we should be able to provide backwards compatibility by copying from the old format to the new format and creating a new array. We can detect the new vs. old ABI by expecting that existing applications will call the __tgt_register_image function. We will create a new __tgt_register_image_v2 function for example that all new programs will call. In libomptarget we then change __tgt_register_image to do the necessary translation.

struct __tgt_bin_desc {                                                                                                                                                                  
  int32_t NumDeviceImages;           // Number of device types supported                                                                                                                   
  __tgt_device_image *DeviceImages;  // Array of device images (1 per dev. type)                                                        
  __tgt_offload_entry *HostEntriesBegin; // Begin of table with all host entries                                                           
  __tgt_offload_entry *HostEntriesEnd;   // End of table (non inclusive)                                                                                                           
};

EXTERN void __tgt_register_lib(__tgt_bin_desc *desc) {
  __tgt_device_image_v2 *new_image = alloc_new_version(desc);
  desc->DeviceImages = new_Image;
  __tgt_register_lib_v2(desc);
}

EXTERN void __tgt_unregister_lib(__tgt_bin_desc *desc) {
  __tgt_unregister_lib_v2(desc);
  dealloc(desc->DeviceImages);
}

Now the rest of libomptarget solely uses the new format, and we check if information is available by seeing that the ImageInfo field is non-null.

[saiislam]

In D124525#3491170, @jhuber6 wrote:

I'm suggesting instead we define a new __tgt_device_image and a new __tgt_register_lib function to support this. This new __tgt_device_image will simply contain a pointer to an optional information struct.

struct __tgt_device_image_v2 {
  void* 	ImageStart;
  void* 	ImageEnd;
  __tgt_offload_entry* 	EntriesBegin;
  __tgt_offload_entry* 	EntriesEnd;
  __tgt_image_into*  ImageInfo;
};

This new struct breaks the ABI with the old __tgt_device_image because these are put into an array and we change the size, but we should be able to provide backwards compatibility by copying from the old format to the new format and creating a new array. We can detect the new vs. old ABI by expecting that existing applications will call the __tgt_register_image function. We will create a new __tgt_register_image_v2 function for example that all new programs will call. In libomptarget we then change __tgt_register_image to do the necessary translation.

struct __tgt_bin_desc {                                                                                                                                                                  
  int32_t NumDeviceImages;           // Number of device types supported                                                                                                                   
  __tgt_device_image *DeviceImages;  // Array of device images (1 per dev. type)                                                        
  __tgt_offload_entry *HostEntriesBegin; // Begin of table with all host entries                                                           
  __tgt_offload_entry *HostEntriesEnd;   // End of table (non inclusive)                                                                                                           
};

EXTERN void __tgt_register_lib(__tgt_bin_desc *desc) {
  __tgt_device_image_v2 *new_image = alloc_new_version(desc);
  desc->DeviceImages = new_Image;
  __tgt_register_lib_v2(desc);
}

EXTERN void __tgt_unregister_lib(__tgt_bin_desc *desc) {
  __tgt_unregister_lib_v2(desc);
  dealloc(desc->DeviceImages);
}

Now the rest of libomptarget solely uses the new format, and we check if information is available by seeing that the ImageInfo field is non-null.

Thanks for the input. I am going to try it.

[saiislam]

Changed the embedding scheme to add ImageInfo field in __tgt_device_image.

[jhuber6]

Where is the code to change the target registration? We need a new initialization runtime call to use and change the old one to reallocate the __tgt_device_image array. Did you work around that some other way?

clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp

613–614

Does this just pass the architecture to the AMD link? Is this related? If not move it to a separate patch and I'll review it.

1069

Is this necessary?

1152

This should be a StringRef, the data is stable. When you create the struct just use this and it will add the null terminator for you.

Constant *ConstantStr = ConstantDataArray::getString(M.getContext(), Str);

1164

An empty StringRef should already create a null string when you create a ConstantDataArray, should be fine to just push it back.

clang/tools/clang-linker-wrapper/OffloadWrapper.cpp

62–63

The struct should be contained within the image itself so we can query it during registration, why do we need to know the number?

64–65

187–190

You should be able to pass the name to the constructor, if it's internal LLVM will give you a new name.

209–243

I would just assert that the metadata vector's size matches the number of images, we're always going to generate these now so it makes sense to generate them in pairs. Then just use llvm::zip or just an iterator to go through both of these at the same time.

210–222

Why do we need all these strings? Should just be generating a constant initializer.

224–227

Why does this need a custom section? We should just use it like this, not sure why we need these to be anything but some internal struct.

for (auto *Image : BinDesc->Images) {
  if (Image->Info)
   // Use Info
}

235–236

Why do we need this when we should already know which image we're looking at?

openmp/libomptarget/include/omptarget.h

122

This comment is extremely long. It should be sufficient to just say what it does, 3 lines max. Specify that it may be null in the device image struct.

145–146

I still feel like we don't need these, correct me if I'm wrong.

161

173–180

Shouldn't the plugin handle this? We should be able to get the architecture in the plugin and select the info struct that matches it, if it exists. Then if the plugin finds one that matches we indicate that it's compatible like we do now. Correct me if I'm wrong.

openmp/libomptarget/src/rtl.cpp

443

Shouldn't we be able to have the plugin handle its own architecture detection, then return from R.is_valid_binary? That makes more sense to me than introducing a new tool to libomptarget, since we can already query the architecture from the CUDA runtime for example.

[saiislam]

Addressed some simple review changes. Will update remaining in the next iteration.

[saiislam]

Thanks for the detailed review. I will update rest of the patch soon.

clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp

613–614

Yes, it passes the architecture to AMD link.

clang/tools/clang-linker-wrapper/OffloadWrapper.cpp

224–227

Custom section is required so that it survives binary stripping done by various OSes. Offload-arch tool will look into this section of the binary to print list of supported architectures.

[jhuber6]

clang/tools/clang-linker-wrapper/OffloadWrapper.cpp

224–227

I see, I was also planning on making a tool that will extract these similar to cuobjdump, we could merge it into that as well.

[jhuber6]

Also I probably should've discussed this earlier, but another potential solution is to use the binary format that we use to embed the object files for the images as well. This is more similar to how CUDA does it. Making it backwards compatible would probably require just assuming it's an image with no information if it doesn't include the magic bytes. This might be better if we want a unified tool to give information on embedded device formats.

[saiislam]

Added the multi-entry logic in libomptarget. Yet to move the image compatibility testing to plugin.

clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp

if (Error Err = createOutputFile(sys::path::filename(ExecutableName) + "-" +
     return std::move(Err);
 
   SmallVector<StringRef, 16> CmdArgs;
   CmdArgs.push_back(*LLDPath);
   CmdArgs.push_back("-flavor");
   CmdArgs.push_back("gnu");
   CmdArgs.push_back("--no-undefined");
   CmdArgs.push_back("-shared");
+  std::string ArchArg = std::string("-plugin-opt=mcpu=").append(Arch.str());
+  CmdArgs.push_back(ArchArg);
   CmdArgs.push_back("-o");
   CmdArgs.push_back(TempFile);
 
   // Add extracted input files.
   for (StringRef Input : InputFiles)
     CmdArgs.push_back(Input);
 
   renderXLinkerArgs(CmdArgs, TheTriple.getTriple());
if (LinkerInputMap.count(File))
       LinkerInputMap[File].push_back(File.Filename);
 
   // Try to link each device toolchain.
   for (auto &LinkerInput : LinkerInputMap) {
     DeviceFile &File = LinkerInput.getFirst();
     Triple TheTriple = Triple(File.TheTriple);
     auto &LinkerInputFiles = LinkerInput.getSecond();
     bool WholeProgram = false;
 
     // Run LTO on any bitcode files and replace the input with the result.
     if (Error Err = linkBitcodeFiles(LinkerInputFiles, TheTriple, File.Arch,
                                      WholeProgram))
       return Err;
 
     if (EmbedBitcode) {
       // If we are embedding bitcode for JIT, skip the final device linking.
       if (LinkerInputFiles.size() != 1 || !WholeProgram)
if (WholeProgram && TheTriple.isNVPTX()) {
                                   LinkerInputFiles.front());
       continue;
     }
 
     auto ImageOrErr = linkDevice(LinkerInputFiles, TheTriple, File.Arch);
     if (!ImageOrErr)
       return ImageOrErr.takeError();
 
     // Create separate images for all the active offload kinds.
     for (OffloadKind Kind : ActiveOffloadKinds[LinkerInput.getFirst()])
       LinkedImages.emplace_back(Kind, TheTriple.getTriple(), File.Arch,
                                 *ImageOrErr);
   }
   return Error::success();
 }
 
 // Compile the module to an object file using the appropriate target machine for
 // the host triple.
 Expected<std::string> compileModule(Module &M) {
   std::string Msg;
Expected<std::string> compileModule(Module &M) {
   return static_cast<std::string>(ObjectFile);
 }
 
 /// Load all of the OpenMP images into a buffer and pass it to the binary
 /// wrapping function to create the registration code in the module \p M.
 Error wrapOpenMPImages(Module &M, ArrayRef<DeviceFile> Images) {
   SmallVector<std::unique_ptr<MemoryBuffer>, 4> SavedBuffers;
   SmallVector<ArrayRef<char>, 4> ImagesToWrap;
+  SmallVector<ArrayRef<char>, 4> OffloadArchs;
+  std::string Arch;
+  OffloadArchs.reserve(Images.size());
   for (const DeviceFile &File : Images) {
     llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> ImageOrError =
         llvm::MemoryBuffer::getFileOrSTDIN(File.Filename);
     if (std::error_code EC = ImageOrError.getError())
       return createFileError(File.Filename, EC);
     ImagesToWrap.emplace_back((*ImageOrError)->getBufferStart(),
                               (*ImageOrError)->getBufferSize());
     SavedBuffers.emplace_back(std::move(*ImageOrError));
+
+    if (Arch.empty()) {
+      Arch = std::string(File.Arch);
+      Arch.push_back('\0');
+      OffloadArchs.emplace_back(Arch.data(), Arch.size());
+    } else {
+      auto curSize = Arch.size();
+      Arch.append(File.Arch);
+      OffloadArchs.emplace_back(&(Arch.at(curSize)), File.Arch.size());
+    }
   }
 
-  if (Error Err = wrapOpenMPBinaries(M, ImagesToWrap))
+  if (Error Err = wrapOpenMPBinaries(M, ImagesToWrap, OffloadArchs))
     return Err;
   return Error::success();
 }
 
 /// Combine all of the CUDA images into a single fatbinary and pass it to the
 /// binary wrapping function to create the registration code in the module \p M.
 Error wrapCudaImages(Module &M, ArrayRef<DeviceFile> Images) {
   SmallVector<StringRef, 4> InputFiles;

clang/tools/clang-linker-wrapper/OffloadWrapper.h

 //===- OffloadWrapper.h --r-------------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_CLANG_LINKER_WRAPPER_OFFLOAD_WRAPPER_H
 #define LLVM_CLANG_TOOLS_CLANG_LINKER_WRAPPER_OFFLOAD_WRAPPER_H
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/IR/Module.h"
 
 /// Wraps the input device images into the module \p M as global symbols and
 /// registers the images with the OpenMP Offloading runtime libomptarget.
-llvm::Error wrapOpenMPBinaries(llvm::Module &M,
-                               llvm::ArrayRef<llvm::ArrayRef<char>> Images);
+llvm::Error
+wrapOpenMPBinaries(llvm::Module &M, llvm::ArrayRef<llvm::ArrayRef<char>> Images,
+                   llvm::ArrayRef<llvm::ArrayRef<char>> OffloadArchs);
 
 /// Wraps the input fatbinary image into the module \p M as global symbols and
 /// registers the images with the CUDA runtime.
 llvm::Error wrapCudaBinary(llvm::Module &M, llvm::ArrayRef<char> Images);
 
 #endif

clang/tools/clang-linker-wrapper/OffloadWrapper.cpp

EntryTy = StructType::create("__tgt_offload_entry", Type::getInt8PtrTy(C),
                                  Type::getInt32Ty(C), Type::getInt32Ty(C));
   return EntryTy;
 }
 
 PointerType *getEntryPtrTy(Module &M) {
   return PointerType::getUnqual(getEntryTy(M));
 }
 
+// This matches the runtime struct definition of __tgt_image_info
+// declared in openmp/libomptarget/include/omptarget.h /
+// struct __tgt_image_info {
+//   int32_t version;
+//   int8_t *offload_arch;
+// };
+StructType *getImageInfoTy(Module &M) {
+  LLVMContext &C = M.getContext();
+  StructType *ImageInfoTy = StructType::getTypeByName(C, "__tgt_image_info");
+  if (!ImageInfoTy)
+    ImageInfoTy = StructType::create("__tgt_image_info", Type::getInt32Ty(C),
+                                     Type::getInt8PtrTy(C));
+  return ImageInfoTy;
+}
+
+PointerType *getImageInfoPtrTy(Module &M) {
+  return PointerType::getUnqual(getImageInfoTy(M));
+}
+
 // struct __tgt_device_image {
 //   void *ImageStart;
 //   void *ImageEnd;
 //   __tgt_offload_entry *EntriesBegin;
 //   __tgt_offload_entry *EntriesEnd;
+//   __tgt_image_info *ImageInfo;
 // };
 StructType *getDeviceImageTy(Module &M) {
   LLVMContext &C = M.getContext();
   StructType *ImageTy = StructType::getTypeByName(C, "__tgt_device_image");
   if (!ImageTy)
     ImageTy = StructType::create("__tgt_device_image", Type::getInt8PtrTy(C),
                                  Type::getInt8PtrTy(C), getEntryPtrTy(M),
-                                 getEntryPtrTy(M));
+                                 getEntryPtrTy(M), getImageInfoPtrTy(M));
   return ImageTy;
 }
 
 PointerType *getDeviceImagePtrTy(Module &M) {
   return PointerType::getUnqual(getDeviceImageTy(M));
 }
 
 // struct __tgt_bin_desc {

 }
 
 PointerType *getBinDescPtrTy(Module &M) {
   return PointerType::getUnqual(getBinDescTy(M));
 }
 
 /// Creates binary descriptor for the given device images. Binary descriptor
 /// is an object that is passed to the offloading runtime at program startup
 /// and it describes all device images available in the executable or shared
 /// library. It is defined as follows
 ///
 /// __attribute__((visibility("hidden")))
 /// extern __tgt_offload_entry *__start_omp_offloading_entries;
 /// __attribute__((visibility("hidden")))
 /// extern __tgt_offload_entry *__stop_omp_offloading_entries;
 ///
 /// static const char Image0[] = { <Bufs.front() contents> };
 ///  ...
 /// static const char ImageN[] = { <Bufs.back() contents> };
 ///
 /// static const __tgt_device_image Images[] = {
 ///   {
 ///     Image0,                            /*ImageStart*/
 ///     Image0 + sizeof(Image0),           /*ImageEnd*/
 ///     __start_omp_offloading_entries,    /*EntriesBegin*/
-///     __stop_omp_offloading_entries      /*EntriesEnd*/
+///     __stop_omp_offloading_entries,     /*EntriesEnd*/
+///     __tgt_image_info                   /*ImageInfo*/
 ///   },
 ///   ...
 ///   {
 ///     ImageN,                            /*ImageStart*/
 ///     ImageN + sizeof(ImageN),           /*ImageEnd*/
 ///     __start_omp_offloading_entries,    /*EntriesBegin*/
-///     __stop_omp_offloading_entries      /*EntriesEnd*/
+///     __stop_omp_offloading_entries,      /*EntriesEnd*/
+///     __tgt_image_info                   /*ImageInfo*/
 ///   }
 /// };
 ///
 /// static const __tgt_bin_desc BinDesc = {
 ///   sizeof(Images) / sizeof(Images[0]),  /*NumDeviceImages*/
 ///   Images,                              /*DeviceImages*/
 ///   __start_omp_offloading_entries,      /*HostEntriesBegin*/
 ///   __stop_omp_offloading_entries        /*HostEntriesEnd*/
 /// };
 ///
 /// Global variable that represents BinDesc is returned.
-GlobalVariable *createBinDesc(Module &M, ArrayRef<ArrayRef<char>> Bufs) {
+GlobalVariable *createBinDesc(Module &M, ArrayRef<ArrayRef<char>> Bufs,
+                              ArrayRef<ArrayRef<char>> OffloadArchs) {
   LLVMContext &C = M.getContext();
   // Create external begin/end symbols for the offload entries table.
   auto *EntriesB = new GlobalVariable(
       M, getEntryTy(M), /*isConstant*/ true, GlobalValue::ExternalLinkage,
       /*Initializer*/ nullptr, "__start_omp_offloading_entries");
   EntriesB->setVisibility(GlobalValue::HiddenVisibility);
   auto *EntriesE = new GlobalVariable(
       M, getEntryTy(M), /*isConstant*/ true, GlobalValue::ExternalLinkage,
auto *DummyEntry = new GlobalVariable(
       M, DummyInit->getType(), true, GlobalVariable::ExternalLinkage, DummyInit,
       "__dummy.omp_offloading.entry");
   DummyEntry->setSection("omp_offloading_entries");
   DummyEntry->setVisibility(GlobalValue::HiddenVisibility);
 
   auto *Zero = ConstantInt::get(getSizeTTy(M), 0u);
   Constant *ZeroZero[] = {Zero, Zero};
 
+  unsigned int ImgCount = 0;
+  std::string OffloadArchBase = "__offload_arch";
+  std::string OffloadImageBase = "offload_image_info";
+
   // Create initializer for the images array.
   SmallVector<Constant *, 4u> ImagesInits;
   ImagesInits.reserve(Bufs.size());
   for (ArrayRef<char> Buf : Bufs) {
     auto *Data = ConstantDataArray::get(C, Buf);
     auto *Image = new GlobalVariable(M, Data->getType(), /*isConstant*/ true,
                                      GlobalVariable::InternalLinkage, Data,
                                      ".omp_offloading.device_image");
     Image->setUnnamedAddr(GlobalValue::UnnamedAddr::Global);
 
     auto *Size = ConstantInt::get(getSizeTTy(M), Buf.size());
     Constant *ZeroSize[] = {Zero, Size};
 
     auto *ImageB =
         ConstantExpr::getGetElementPtr(Image->getValueType(), Image, ZeroZero);
     auto *ImageE =
         ConstantExpr::getGetElementPtr(Image->getValueType(), Image, ZeroSize);
 
-    ImagesInits.push_back(ConstantStruct::get(getDeviceImageTy(M), ImageB,
-                                              ImageE, EntriesB, EntriesE));
+    auto OArch = OffloadArchs[ImgCount];
+    Constant *OArchV = ConstantDataArray::get(C, OArch);
+    std::string OffloadArchGV(OffloadArchBase),
+        OffloadImageGV(OffloadImageBase);
+    if (ImgCount) {
+      auto Suffix = std::to_string(ImgCount);
+      OffloadArchGV.append(".").append(Suffix);
+      OffloadImageGV.append(".").append(Suffix);
+    }
+
+    auto *GV =
+        new GlobalVariable(M, OArchV->getType(), /*isConstant*/ true,
+                           GlobalValue::InternalLinkage, OArchV, OffloadArchGV);
+    GV->setUnnamedAddr(GlobalValue::UnnamedAddr::Global);
+
+    // store value of these variables (i.e. offload archs) into a custom
+    // section which will be used by "offload-arch -f". It won't be
+    // removed during binary stripping.
+    GV->setSection(".offload_arch_list");
+
+    auto *RequirementVPtr =
+        ConstantExpr::getGetElementPtr(GV->getValueType(), GV, Zero);
+    RequirementVPtr =
+        ConstantExpr::getBitCast(RequirementVPtr, Type::getInt8PtrTy(C));
+    auto *InfoInit = ConstantStruct::get(
+        getImageInfoTy(M), ConstantInt::get(Type::getInt32Ty(C), 1),
+        RequirementVPtr);
+    ++ImgCount;
+    auto *ImageInfoGV =
+        new GlobalVariable(M, InfoInit->getType(),
+                           /*isConstant*/ true, GlobalValue::InternalLinkage,
+                           InfoInit, OffloadImageGV);
+    ImageInfoGV->setUnnamedAddr(GlobalValue::UnnamedAddr::Global);
+    ImagesInits.push_back(ConstantStruct::get(
+        getDeviceImageTy(M), ImageB, ImageE, EntriesB, EntriesE, ImageInfoGV));
   }
 
   // Then create images array.
   auto *ImagesData = ConstantArray::get(
       ArrayType::get(getDeviceImageTy(M), ImagesInits.size()), ImagesInits);
 
   auto *Images =
       new GlobalVariable(M, ImagesData->getType(), /*isConstant*/ true,
void createRegisterFunction(Module &M, GlobalVariable *BinDesc) {
   auto *Func = Function::Create(FuncTy, GlobalValue::InternalLinkage,
                                 ".omp_offloading.descriptor_reg", &M);
   Func->setSection(".text.startup");
 
   // Get __tgt_register_lib function declaration.
   auto *RegFuncTy = FunctionType::get(Type::getVoidTy(C), getBinDescPtrTy(M),
                                       /*isVarArg*/ false);
   FunctionCallee RegFuncC =
-      M.getOrInsertFunction("__tgt_register_lib", RegFuncTy);
+      M.getOrInsertFunction("__tgt_register_lib_v2", RegFuncTy);
 
   // Construct function body
   IRBuilder<> Builder(BasicBlock::Create(C, "entry", Func));
   Builder.CreateCall(RegFuncC, BinDesc);
   Builder.CreateRetVoid();
 
   // Add this function to constructors.
   // Set priority to 1 so that __tgt_register_lib is executed AFTER
   // __tgt_register_requires (we want to know what requirements have been
   // asked for before we load a libomptarget plugin so that by the time the
   // plugin is loaded it can report how many devices there are which can
   // satisfy these requirements).
void createRegisterFatbinFunction(Module &M, GlobalVariable *FatbinDesc) {
   DtorBuilder.CreateRetVoid();
 
   // Add this function to constructors.
   appendToGlobalCtors(M, CtorFunc, /*Priority*/ 1);
 }
 
 } // namespace
 
-Error wrapOpenMPBinaries(Module &M, ArrayRef<ArrayRef<char>> Images) {
-  GlobalVariable *Desc = createBinDesc(M, Images);
+Error wrapOpenMPBinaries(Module &M, ArrayRef<ArrayRef<char>> Images,
+                         ArrayRef<ArrayRef<char>> OffloadArchs) {
+  GlobalVariable *Desc = createBinDesc(M, Images, OffloadArchs);
   if (!Desc)
     return createStringError(inconvertibleErrorCode(),
                              "No binary descriptors created.");
   createRegisterFunction(M, Desc);
   createUnregisterFunction(M, Desc);
   return Error::success();
 }
 

clang/tools/clang-offload-wrapper/ClangOffloadWrapper.cpp

void createRegisterFunction(GlobalVariable *BinDesc) {
     auto *Func = Function::Create(FuncTy, GlobalValue::InternalLinkage,
                                   ".omp_offloading.descriptor_reg", &M);
     Func->setSection(".text.startup");
 
     // Get __tgt_register_lib function declaration.
     auto *RegFuncTy = FunctionType::get(Type::getVoidTy(C), getBinDescPtrTy(),
                                         /*isVarArg*/ false);
     FunctionCallee RegFuncC =
-        M.getOrInsertFunction("__tgt_register_lib", RegFuncTy);
+        M.getOrInsertFunction("__tgt_register_lib_v2", RegFuncTy);
 
     // Construct function body
     IRBuilder<> Builder(BasicBlock::Create(C, "entry", Func));
     Builder.CreateCall(RegFuncC, BinDesc);
     Builder.CreateRetVoid();
 
     // Add this function to constructors.
     // Set priority to 1 so that __tgt_register_lib is executed AFTER

openmp/libomptarget/include/omptarget.h

 struct __tgt_offload_entry {
   void *addr;   // Pointer to the offload entry info (function or global)
   char *name;   // Name of the function or global
   size_t size;  // Size of the entry info (0 if it is a function)
   int32_t flags; // Flags associated with the entry, e.g. 'link'.
   int32_t reserved; // Reserved, to be used by the runtime library.
 };
 
+/// __tgt_image_info:
+///
+/// The information in this struct is provided in the clang-linker-wrapper
+/// as a call to __tgt_register_image_info for each image in the library
+/// of images also created by the clang-linker-wrapper.
+/// __tgt_register_image_info is called for each image BEFORE the single
+/// call to __tgt_register_lib so that image information is available
+/// before they are loaded. clang-linker-wrapper gets this image information
+/// from command line arguments provided by the clang driver when it creates
+/// the call to the __clang-linker-wrapper command.
+/// This architecture allows the binary image (pointed to by ImageStart and
+/// ImageEnd in __tgt_device_image) to remain architecture indenendent.
+/// That is, the architecture independent part of the libomptarget runtime
+/// does not need to peer inside the image to determine if it is loadable
+/// even though in most cases the image is an elf object.
+/// There is one __tgt_image_info for each __tgt_device_image. For backward
+/// compabibility, no changes are allowed to either __tgt_device_image or
+/// __tgt_bin_desc. The absense of __tgt_image_info is the indication that
+/// the runtime is being used on a binary created by an old version of
+/// the compiler.
+///
+struct __tgt_image_info {
+  int32_t image_info_version;           // The version of this struct (= 1)
+  char *offload_arch;        // e.g. sm_30, sm_70, gfx906, includes features
+};
+
 /// This struct is a record of the device image information
+struct __tgt_device_image_old {
+  void *ImageStart;                  // Pointer to the target code start
+  void *ImageEnd;                    // Pointer to the target code end
+  __tgt_offload_entry *EntriesBegin; // Begin of table with all target entries
+  __tgt_offload_entry *EntriesEnd;   // End of table (non inclusive)
+};
+
 struct __tgt_device_image {
   void *ImageStart;                  // Pointer to the target code start
   void *ImageEnd;                    // Pointer to the target code end
   __tgt_offload_entry *EntriesBegin; // Begin of table with all target entries
   __tgt_offload_entry *EntriesEnd;   // End of table (non inclusive)
+  __tgt_image_info *ImageInfo;       // Metadata about the image (may be null).
 };
 
 /// This struct is a record of all the host code that may be offloaded to a
 /// target.
 struct __tgt_bin_desc {
   int32_t NumDeviceImages;           // Number of device types supported
   __tgt_device_image *DeviceImages;  // Array of device images (1 per dev. type)
   __tgt_offload_entry *HostEntriesBegin; // Begin of table with all host entries
   __tgt_offload_entry *HostEntriesEnd;   // End of table (non inclusive)
 };
 
+/// __tgt_active_offload_env
+///
+/// This structure is created by __tgt_get_active_offload_env and is used
+/// to determine compatibility of the images with the current environment
+/// that is "in play".
+struct __tgt_active_offload_env {
+  char *capabilities; // string returned by offload-arch -c
+};
+
 /// This struct contains the offload entries identified by the target runtime
 struct __tgt_target_table {
   __tgt_offload_entry *EntriesBegin; // Begin of the table with all the entries
   __tgt_offload_entry
       *EntriesEnd; // End of the table with all the entries (non inclusive)
 };
 
 // clang-format on
 
 /// This struct contains information exchanged between different asynchronous
 /// operations for device-dependent optimization and potential synchronization
 struct __tgt_async_info {
   // A pointer to a queue-like structure where offloading operations are issued.
   // We assume to use this structure to do synchronization. In CUDA backend, it
   // is CUstream.
   void *Queue = nullptr;
 };
 
 struct DeviceTy;
 
 /// The libomptarget wrapper around a __tgt_async_info object directly
 /// associated with a libomptarget layer device. RAII semantics to avoid
 /// mistakes.
 class AsyncInfoTy {

 /// Dummy target so we have a symbol for generating host fallback.
 void *llvm_omp_target_dynamic_shared_alloc();
 
 /// add the clauses of the requires directives in a given file
 void __tgt_register_requires(int64_t flags);
 
 /// adds a target shared library to the target execution image
 void __tgt_register_lib(__tgt_bin_desc *desc);
+void __tgt_register_lib_v2(__tgt_bin_desc *desc);
 
 /// Initialize all RTLs at once
 void __tgt_init_all_rtls();
 
 /// removes a target shared library from the target execution image
 void __tgt_unregister_lib(__tgt_bin_desc *desc);
 
 // creates the host to target data mapping, stores it in the

openmp/libomptarget/include/rtl.h

struct RTLsTy {
   // Initialize RTL if it has not been initialized
   void initRTLonce(RTLInfoTy &RTL);
 
   // Initialize all RTLs
   void initAllRTLs();
 
   // Register a shared library with all (compatible) RTLs.
   void RegisterLib(__tgt_bin_desc *desc);
+  void RegisterLibV2(__tgt_bin_desc *desc);
 
   // Unregister a shared library from all RTLs.
   void UnregisterLib(__tgt_bin_desc *desc);
 
   // Mutex-like object to guarantee thread-safety and unique initialization
   // (i.e. the library attempts to load the RTLs (plugins) only once).
   std::once_flag initFlag;
   void LoadRTLs(); // not thread-safe

openmp/libomptarget/src/exports

 VERS1.0 {
   global:
     __tgt_register_requires;
     __tgt_register_lib;
+    __tgt_register_lib_v2;
     __tgt_unregister_lib;
     __tgt_init_all_rtls;
     __tgt_target_data_begin;
     __tgt_target_data_end;
     __tgt_target_data_update;
     __tgt_target;
     __tgt_target_teams;
     __tgt_target_data_begin_nowait;

openmp/libomptarget/src/interface.cpp

if (RTL.register_lib) {
       if ((*RTL.register_lib)(desc) != OFFLOAD_SUCCESS) {
         DP("Could not register library with %s", RTL.RTLName.c_str());
       }
     }
   }
   PM->RTLs.RegisterLib(desc);
 }
 
+EXTERN void __tgt_register_lib_v2(__tgt_bin_desc *desc) {
+  TIMESCOPE();
+  std::call_once(PM->RTLs.initFlag, &RTLsTy::LoadRTLs, &PM->RTLs);
+  for (auto &RTL : PM->RTLs.AllRTLs) {
+    if (RTL.register_lib) {
+      if ((*RTL.register_lib)(desc) != OFFLOAD_SUCCESS) {
+        DP("Could not register library with %s", RTL.RTLName.c_str());
+      }
+    }
+  }
+  PM->RTLs.RegisterLibV2(desc);
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 /// Initialize all available devices without registering any image
 EXTERN void __tgt_init_all_rtls() { PM->RTLs.initAllRTLs(); }
 
 ////////////////////////////////////////////////////////////////////////////////
 /// unloads a target shared library
 EXTERN void __tgt_unregister_lib(__tgt_bin_desc *desc) {
   TIMESCOPE();

openmp/libomptarget/src/rtl.cpp

 //===----------- rtl.cpp - Target independent OpenMP target RTL -----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Functionality for handling RTL plugins.
 //
 //===----------------------------------------------------------------------===//
 
 #include "rtl.h"
 #include "device.h"
 #include "private.h"
+//#include "llvm/OffloadArch/OffloadArch.h"
 
 #include <cassert>
 #include <cstdlib>
 #include <cstring>
 #include <dlfcn.h>
 #include <mutex>
 #include <string>
+// It's strange we do not have llvm tools for openmp runtime, so we use stat
+#include <sys/stat.h>
 
 // List of all plugins that can support offloading.
 static const char *RTLNames[] = {
     /* PowerPC target       */ "libomptarget.rtl.ppc64.so",
     /* x86_64 target        */ "libomptarget.rtl.x86_64.so",
     /* CUDA target          */ "libomptarget.rtl.cuda.so",
     /* AArch64 target       */ "libomptarget.rtl.aarch64.so",
     /* SX-Aurora VE target  */ "libomptarget.rtl.ve.so",
void RTLsTy::initRTLonce(RTLInfoTy &R) {
   }
 }
 
 void RTLsTy::initAllRTLs() {
   for (auto &R : AllRTLs)
     initRTLonce(R);
 }
 
+/// Query runtime capabilities of this system by calling offload-arch -c
+/// offload_arch_output_buffer is persistant storage returned by this
+/// __tgt_get_active_offload_env.
+static void
+__tgt_get_active_offload_env(__tgt_active_offload_env *active_env,
+                             char *offload_arch_output_buffer,
+                             size_t offload_arch_output_buffer_size) {
+
+  // If OFFLOAD_ARCH_OVERRIDE env varible is present then use its value instead
+  // of querying it using LLVMOffloadArch library.
+  if (char *OffloadArchEnvVar = getenv("OFFLOAD_ARCH_OVERRIDE")) {
+    if (OffloadArchEnvVar) {
+      active_env->capabilities = OffloadArchEnvVar;
+      return;
+    }
+  }
+  // Qget runtime capabilities of this system with libLLVMOffloadArch.a
+  // if (int rc = getRuntimeCapabilities(offload_arch_output_buffer,
+  //                                     offload_arch_output_buffer_size))
+  //   return;
+  // active_env->capabilities = offload_arch_output_buffer;
+  // return;
+}
+
+std::vector<std::string> _splitstrings(char *input, const char *sep) {
+  std::vector<std::string> split_strings;
+  std::string s(input);
+  std::string delimiter(sep);
+  size_t pos = 0;
+  while ((pos = s.find(delimiter)) != std::string::npos) {
+    if (pos != 0)
+      split_strings.push_back(s.substr(0, pos));
+    s.erase(0, pos + delimiter.length());
+  }
+  if (s.length() > 1)
+    split_strings.push_back(s.substr(0, s.length()));
+  return split_strings;
+}
+
+static bool _ImageIsCompatibleWithEnv(__tgt_image_info *image_info,
+                                      __tgt_active_offload_env *active_env) {
+  // get_image_info will return null if no image information was registered.
+  // If no image information, assume application built with old compiler and
+  // check each image.
+  if (!(image_info && image_info->image_info_version == 1))
+    return true;
+
+  if (!active_env->capabilities)
+    return false;
+
+  // Each runtime requirement for the compiled image is stored in
+  // the image_info->offload_arch (TargetID) string.
+  // Each runtime capability obtained from "offload-arch -c" is stored in
+  // actvie_env->capabilities (TargetID) string.
+  // If every requirement has a matching capability, then the image
+  // is compatible with active environment
+
+  std::vector<std::string> reqs = _splitstrings(image_info->offload_arch, ":");
+  std::vector<std::string> caps = _splitstrings(active_env->capabilities, ":");
+
+  bool is_compatible = true;
+  for (auto req : reqs) {
+    bool missing_capability = true;
+    for (auto capability : caps)
+      if (capability == req)
+        missing_capability = false;
+    if (missing_capability) {
+      DP("Image requires %s but runtime capability %s is missing.\n",
+         image_info->offload_arch, req.c_str());
+      is_compatible = false;
+    }
+  }
+  return is_compatible;
+}
+
 void RTLsTy::RegisterLib(__tgt_bin_desc *desc) {
+
+  __tgt_device_image *newDeviceImages =
+      new __tgt_device_image[desc->NumDeviceImages];
+
+  for (int32_t i = 0; i < desc->NumDeviceImages; ++i) {
+    newDeviceImages[i].EntriesBegin = desc->DeviceImages[i].EntriesBegin;
+    newDeviceImages[i].EntriesEnd = desc->DeviceImages[i].EntriesEnd;
+    newDeviceImages[i].ImageStart = desc->DeviceImages[i].ImageStart;
+    newDeviceImages[i].ImageEnd = desc->DeviceImages[i].ImageEnd;
+    newDeviceImages[i].ImageInfo = nullptr;
+    // TODO : delete(desc->DeviceImages[i]);
+  }
+
+  desc->DeviceImages = static_cast<__tgt_device_image *>(newDeviceImages);
+
+  this->RegisterLibV2(desc);
+}
+
+#define MAX_CAPS_STR_SIZE 1024
+void RTLsTy::RegisterLibV2(__tgt_bin_desc *desc) {
+
+  // Get the current active offload environment
+  __tgt_active_offload_env offload_env = {nullptr};
+  // Need a buffer to hold results of offload-arch -c command
+  size_t offload_arch_output_buffer_size = MAX_CAPS_STR_SIZE;
+  std::vector<char> offload_arch_output_buffer;
+  offload_arch_output_buffer.resize(offload_arch_output_buffer_size);
+  __tgt_get_active_offload_env(&offload_env, offload_arch_output_buffer.data(),
+                               offload_arch_output_buffer_size);
+
+  RTLInfoTy *FoundRTL = NULL;
   PM->RTLsMtx.lock();
   // Register the images with the RTLs that understand them, if any.
   for (int32_t i = 0; i < desc->NumDeviceImages; ++i) {
     // Obtain the image.
     __tgt_device_image *img = &desc->DeviceImages[i];
 
-    RTLInfoTy *FoundRTL = nullptr;
-
+    // Get corresponding image info offload_arch and check with runtime
+    if (!_ImageIsCompatibleWithEnv(img->ImageInfo, &offload_env))
+      continue;
+    FoundRTL = NULL;
     // Scan the RTLs that have associated images until we find one that supports
     // the current image.
     for (auto &R : AllRTLs) {
+
       if (!R.is_valid_binary(img)) {
         DP("Image " DPxMOD " is NOT compatible with RTL %s!\n",
            DPxPTR(img->ImageStart), R.RTLName.c_str());
         continue;
       }
 
       DP("Image " DPxMOD " is compatible with RTL %s!\n",
          DPxPTR(img->ImageStart), R.RTLName.c_str());
for (int32_t i = 0; i < desc->NumDeviceImages; ++i) {
     }
 
     if (!FoundRTL) {
       DP("No RTL found for image " DPxMOD "!\n", DPxPTR(img->ImageStart));
     }
   }
   PM->RTLsMtx.unlock();
 
+  if (!FoundRTL) {
+    if (PM->TargetOffloadPolicy == tgt_mandatory)
+      fprintf(stderr, "ERROR:\
+  Runtime capabilities do NOT meet any offload image offload_arch\n\
+  and the OMP_TARGET_OFFLOAD policy is mandatory.  Terminating!\n\
+  Runtime capabilities : %s\n",
+              offload_env.capabilities);
+    else if (PM->TargetOffloadPolicy == tgt_disabled)
+      fprintf(stderr, "WARNING: Offloading is disabled.\n");
+    else
+      fprintf(
+          stderr,
+          "WARNING: Runtime capabilities do NOT meet any image offload_arch.\n\
+   So device offloading is now disabled.\n\
+  Runtime capabilities : %s\n",
+          offload_env.capabilities);
+    if (PM->TargetOffloadPolicy != tgt_disabled) {
+      for (int32_t i = 0; i < desc->NumDeviceImages; ++i) {
+        __tgt_image_info *image_info = desc->DeviceImages[i].ImageInfo;
+        if (image_info && image_info->image_info_version == 1)
+          fprintf(stderr, "\
+    Image %d offload_arch : %s\n",
+                  i, image_info->offload_arch);
+        else
+          fprintf(stderr, "\
+    Image %d has no offload_arch. Could be from older compiler\n",
+                  i);
+      }
+    }
+    if (PM->TargetOffloadPolicy == tgt_mandatory)
+      exit(1);
+  }
+
   DP("Done registering entries!\n");
 }
 
 void RTLsTy::UnregisterLib(__tgt_bin_desc *desc) {
   DP("Unloading target library!\n");
 
   PM->RTLsMtx.lock();
   // Find which RTL understands each image, if any.