[Cuda] Add initial support for wrapping CUDA images in the new driver.

[jhuber6]

Summary

This patch adds the initial support for wrapping CUDA images. This
requires changing some of the logic for how we bundle images. We now
need to copy the image for all kinds that are active for the
architecture. Then we need to run a separate wrapping job if the Kind is
Cuda. For cuda wrapping we need to use the fatbinary program from the
CUDA SDK to bundle all the binaries together. This is then passed to a
new function to perfom the actual module code generation that will be
implemented in a later patch.

Depends on D120273 D123471

[jhuber6]

rebase.

[jhuber6]

Updating to use the OffloadKind enum rather than the string. I will also probably simplify some of the logic for handling multiple files here in a later patch.

[jhuber6]

Fix test

[tra]

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

188

Extend enum OffloadKind to include these special kinds.

201

Why do we limit ourselves to uint16_t here? Can't we just use OffloadKind itself and get rid of these casts?

726

We should have a way to pass extra options to fatbinary, too. E.g. we may want to use --compress-all. Also, we may need to pass through -g for debug builds.

Oh. Debug builds. Makes me wonder if cuda-gdb will be able to find GPU binaries packaged by the new driver. If it does not, it will be a rather serious problem. It would likely affect various profiling tools the same way, too. Can you give it a try?

1192–1193

Nit. We should think of changing linkBitcodeFiles to return llvm::ErrorOr<bool> so we can return WholeArchive value, instead of modifying it as an argument.

1200

What's expected to happen if we have more than one input?
If only one is ever expected, I'd add an assert.

1206–1207

Should EmbedBitcode be mutually exclusive vs WholeArchive?

Can we ever end up with both unset?

1210

ditto about the assert on the number of inputs.

1347

The M contains generated registration glue at this point, right? It may be worth a comment to explain what is it that we're compiling here.

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

271

A "not-implemented-yet/TODO" comment would be appropriate here.

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

20

It should be either "Wraps/registers" or "Wrap/register".

[jhuber6]

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

726

I was planning on implementing this stuff more generally when we get the new binary tool in D125165 landed. That will allow me to more generally put any number of command line arguments into the binary itself and fish it out here. We already support a janky version for the -Xcuda-ptxas option here, but it's a mess and I'm planning on getting rid of it. Is it okay to punt that into the future? Debug builds are another sore point I don't handle super well right now but will be addressed better with D125165.

I haven't tested cuda-gdb, but I embed the fatbinary the same way that we do in non-rdc mode. I can read them with cuobjdump in the final executable so I'm assuming it's compatible.

1192–1193

That's a good idea, I'm planning on cleaning a lot of this stuff up later.

1200

This isn't used right now since JIT hasn't made it in. It should probably be a proper error honestly.

1206–1207

EmbedBitcode might need to require WholeArchive considering that it's supposed to be a completely linked image that can be sent to a JIT engine.

1347

Sure.

[jhuber6]

Addressing comments

[tra]

Please add a TODO around the items that need further work, so we don't forget about them. Review comments tend to fade from memory.

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

191

Is there a particular reason for multiplying by 37? Enum values by themselves should do the job just fine.

726

I can read them with cuobjdump in the final executable so I'm assuming it's compatible.

We should be OK then.

726

Debug builds are another sore point I don't handle super well right now but will be addressed better with D125165

OK.

[jhuber6]

Thanks for the review.

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

191

That's what LLVM does for the regular DenseMapInfo<uint16_t>::getHashValue() so I just copied it here.

[tra]

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

191

It would make sense for mapping the full range of uint16_t into a much smaller set of entries. In this case we're already dealing with a very small densely packed set of values. For all practical purposes is a convenient overkill We could get by with just using a vector+direct indexing. We also don't care about hash collisions even if they happen.

Removing multiplication would not make much of a difference, but it would be one less question for the reader to ask, when they look at this code.

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

 #define OFFLOAD_SECTION_MAGIC_STR ".llvm.offloading"
 
 /// The magic offset for the first object inside CUDA's fatbinary. This can be
 /// different but it should work for what is passed here.
 static constexpr unsigned FatbinaryOffset = 0x50;
 
 /// Information for a device offloading file extracted from the host.
 struct DeviceFile {
-  DeviceFile(StringRef Kind, StringRef TheTriple, StringRef Arch,
+  DeviceFile(OffloadKind Kind, StringRef TheTriple, StringRef Arch,
              StringRef Filename)
       : Kind(Kind), TheTriple(TheTriple), Arch(Arch), Filename(Filename) {}
 
-  std::string Kind;
+  OffloadKind Kind;
   std::string TheTriple;
   std::string Arch;
   std::string Filename;
 };
 
 namespace llvm {
 /// Helper that allows DeviceFile to be used as a key in a DenseMap. For now we
 /// assume device files with matching architectures and triples but different
 /// offloading kinds should be handlded together, this may not be true in the
 /// future.
+
+// Provide DenseMapInfo for OffloadKind.
+template <> struct DenseMapInfo<OffloadKind> {
+  static inline OffloadKind getEmptyKey() {
+    return static_cast<OffloadKind>(0xFFFF);
+  }
+  static inline OffloadKind getTombstoneKey() {
+    return static_cast<OffloadKind>(0xFFFF - 1);
+  }
+  static unsigned getHashValue(const OffloadKind &Val) { return Val * 37U; }
+
+  static bool isEqual(const OffloadKind &LHS, const OffloadKind &RHS) {
+    return LHS == RHS;
+  }
+};
 template <> struct DenseMapInfo<DeviceFile> {
   static DeviceFile getEmptyKey() {
-    return {DenseMapInfo<StringRef>::getEmptyKey(),
+    return {static_cast<OffloadKind>(DenseMapInfo<uint16_t>::getEmptyKey()),
             DenseMapInfo<StringRef>::getEmptyKey(),
             DenseMapInfo<StringRef>::getEmptyKey(),
             DenseMapInfo<StringRef>::getEmptyKey()};
   }
   static DeviceFile getTombstoneKey() {
-    return {DenseMapInfo<StringRef>::getTombstoneKey(),
+    return {static_cast<OffloadKind>(DenseMapInfo<uint16_t>::getTombstoneKey()),
             DenseMapInfo<StringRef>::getTombstoneKey(),
             DenseMapInfo<StringRef>::getTombstoneKey(),
             DenseMapInfo<StringRef>::getTombstoneKey()};
   }
   static unsigned getHashValue(const DeviceFile &I) {
     return DenseMapInfo<StringRef>::getHashValue(I.TheTriple) ^
            DenseMapInfo<StringRef>::getHashValue(I.Arch);
   }
std::string getMainExecutable(const char *Name) {
   return sys::path::parent_path(COWPath).str();
 }
 
 /// Extract the device file from the string '<kind>-<triple>-<arch>=<library>'.
 DeviceFile getBitcodeLibrary(StringRef LibraryStr) {
   auto DeviceAndPath = StringRef(LibraryStr).split('=');
   auto StringAndArch = DeviceAndPath.first.rsplit('-');
   auto KindAndTriple = StringAndArch.first.split('-');
-  return DeviceFile(KindAndTriple.first, KindAndTriple.second,
+  return DeviceFile(getOffloadKind(KindAndTriple.first), KindAndTriple.second,
                     StringAndArch.second, DeviceAndPath.second);
 }
 
 /// Get a temporary filename suitable for output.
 Error createOutputFile(const Twine &Prefix, StringRef Extension,
                        SmallString<128> &NewFilename) {
   if (!SaveTemps) {
     if (std::error_code EC =
while (Offset < Contents.size()) {
     if (!OutputOrErr)
       return OutputOrErr.takeError();
     std::unique_ptr<FileOutputBuffer> Output = std::move(*OutputOrErr);
     std::copy(Binary.getImage().bytes_begin(), Binary.getImage().bytes_end(),
               Output->getBufferStart());
     if (Error E = Output->commit())
       return E;
 
-    DeviceFiles.emplace_back(Kind, Binary.getTriple(), Binary.getArch(),
-                             TempFile);
+    DeviceFiles.emplace_back(Binary.getOffloadKind(), Binary.getTriple(),
+                             Binary.getArch(), TempFile);
 
     Offset += Binary.getSize();
   }
 
   return Error::success();
 }
 
 Expected<Optional<std::string>>
Expected<std::string> link(ArrayRef<std::string> InputFiles, Triple TheTriple,
   for (StringRef Input : InputFiles)
     CmdArgs.push_back(Input);
 
   if (Error Err = executeCommands(*NvlinkPath, CmdArgs))
     return std::move(Err);
 
   return static_cast<std::string>(TempFile);
 }
+
+Expected<std::string> fatbinary(ArrayRef<StringRef> InputFiles,
+                                Triple TheTriple, ArrayRef<StringRef> Archs) {
+  // NVPTX uses the fatbinary program to bundle the linked images.
+  Expected<std::string> FatBinaryPath =
+      findProgram("fatbinary", {CudaBinaryPath});
+  if (!FatBinaryPath)
+    return FatBinaryPath.takeError();
+
+  // Create a new file to write the linked device image to.
+  SmallString<128> TempFile;
+  if (Error Err = createOutputFile(sys::path::filename(ExecutableName) +
+                                       "-device-" + TheTriple.getArchName(),
+                                   "fatbin", TempFile))
+    return std::move(Err);
+
+  BumpPtrAllocator Alloc;
+  StringSaver Saver(Alloc);
+
+  SmallVector<StringRef, 16> CmdArgs;
+  CmdArgs.push_back(*FatBinaryPath);
+  CmdArgs.push_back(TheTriple.isArch64Bit() ? "-64" : "-32");
+  CmdArgs.push_back("--create");
+  CmdArgs.push_back(TempFile);
+  for (const auto &FileAndArch : llvm::zip(InputFiles, Archs))
+    CmdArgs.push_back(Saver.save("--image=profile=" + std::get<1>(FileAndArch) +
+                                 ",file=" + std::get<0>(FileAndArch)));
+
+  if (Error Err = executeCommands(*FatBinaryPath, CmdArgs))
+    return std::move(Err);
+
+  return static_cast<std::string>(TempFile);
+}
 } // namespace nvptx
 namespace amdgcn {
 Expected<std::string> link(ArrayRef<std::string> InputFiles, Triple TheTriple,
                            StringRef Arch) {
   // AMDGPU uses lld to link device object files.
   Expected<std::string> LLDPath =
       findProgram("lld", {getMainExecutable("lld")});
   if (!LLDPath)
Error linkBitcodeFiles(SmallVectorImpl<std::string> &InputFiles,
   InputFiles = NewInputFiles;
 
   return Error::success();
 }
 
 /// Runs the appropriate linking action on all the device files specified in \p
 /// DeviceFiles. The linked device images are returned in \p LinkedImages.
 Error linkDeviceFiles(ArrayRef<DeviceFile> DeviceFiles,
-                      SmallVectorImpl<std::string> &LinkedImages) {
-  // Get the list of inputs for a specific device.
+                      SmallVectorImpl<DeviceFile> &LinkedImages) {
+  // Get the list of inputs and active offload kinds for a specific device.
   DenseMap<DeviceFile, SmallVector<std::string, 4>> LinkerInputMap;
-  for (auto &File : DeviceFiles)
+  DenseMap<DeviceFile, DenseSet<OffloadKind>> ActiveOffloadKinds;
+  for (auto &File : DeviceFiles) {
+    ActiveOffloadKinds[File].insert(File.Kind);
     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(LinkerInput.getSecond(), TheTriple,
-                                     File.Arch, WholeProgram))
+    if (Error Err = linkBitcodeFiles(LinkerInputFiles, TheTriple, File.Arch,
+                                     WholeProgram))
       return Err;
 
     // If we are embedding bitcode for JIT, skip the final device linking.
     if (EmbedBitcode) {
-      assert(!LinkerInput.getSecond().empty() && "No bitcode image to embed");
-      LinkedImages.push_back(LinkerInput.getSecond().front());
+      assert(!LinkerInputFiles.empty() && "No bitcode image to embed");
+      LinkedImages.emplace_back(OFK_OpenMP, TheTriple.getTriple(), File.Arch,
+                                LinkerInputFiles.front());
       continue;
     }
 
     // If we performed LTO on NVPTX and had whole program visibility, we can use
     // CUDA in non-RDC mode.
     if (WholeProgram && TheTriple.isNVPTX()) {
-      assert(!LinkerInput.getSecond().empty() && "No non-RDC image to embed");
-      LinkedImages.push_back(LinkerInput.getSecond().front());
+      assert(!LinkerInputFiles.empty() && "No non-RDC image to embed");
+      for (OffloadKind Kind : ActiveOffloadKinds[LinkerInput.getFirst()])
+        LinkedImages.emplace_back(Kind, TheTriple.getTriple(), File.Arch,
+                                  LinkerInputFiles.front());
       continue;
     }
 
-    auto ImageOrErr = linkDevice(LinkerInput.getSecond(), TheTriple, File.Arch);
+    auto ImageOrErr = linkDevice(LinkerInputFiles, TheTriple, File.Arch);
     if (!ImageOrErr)
       return ImageOrErr.takeError();
 
-    LinkedImages.push_back(*ImageOrErr);
+    // 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) {
   if (TM->addPassesToEmitFile(CodeGenPasses, *OS, nullptr, CGFT_ObjectFile))
     return createStringError(inconvertibleErrorCode(),
                              "Failed to execute host backend");
   CodeGenPasses.run(M);
 
   return static_cast<std::string>(ObjectFile);
 }
 
-/// Creates the object file containing the device image and runtime registration
-/// code from the device images stored in \p Images.
-Expected<std::string> wrapDeviceImages(ArrayRef<std::string> Images) {
+/// 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;
-
-  for (StringRef ImageFilename : Images) {
+  for (const DeviceFile &File : Images) {
     llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> ImageOrError =
-        llvm::MemoryBuffer::getFileOrSTDIN(ImageFilename);
+        llvm::MemoryBuffer::getFileOrSTDIN(File.Filename);
     if (std::error_code EC = ImageOrError.getError())
-      return createFileError(ImageFilename, EC);
+      return createFileError(File.Filename, EC);
     ImagesToWrap.emplace_back((*ImageOrError)->getBufferStart(),
                               (*ImageOrError)->getBufferSize());
     SavedBuffers.emplace_back(std::move(*ImageOrError));
   }
 
+  if (Error Err = wrapOpenMPBinaries(M, ImagesToWrap))
+    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;
+  SmallVector<StringRef, 4> Architectures;
+  for (const DeviceFile &File : Images) {
+    InputFiles.push_back(File.Filename);
+    Architectures.push_back(File.Arch);
+  }
+
+  // CUDA expects its embedded device images to be a fatbinary.
+  Triple TheTriple = Triple(Images.front().TheTriple);
+  auto FileOrErr = nvptx::fatbinary(InputFiles, TheTriple, Architectures);
+  if (!FileOrErr)
+    return FileOrErr.takeError();
+
+  llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> ImageOrError =
+      llvm::MemoryBuffer::getFileOrSTDIN(*FileOrErr);
+  if (std::error_code EC = ImageOrError.getError())
+    return createFileError(*FileOrErr, EC);
+
+  auto ImageToWrap = ArrayRef<char>((*ImageOrError)->getBufferStart(),
+                                    (*ImageOrError)->getBufferSize());
+
+  if (Error Err = wrapCudaBinary(M, ImageToWrap))
+    return Err;
+  return Error::success();
+}
+
+/// Creates the object file containing the device image and runtime
+/// registration code from the device images stored in \p Images.
+Expected<SmallVector<std::string, 2>>
+wrapDeviceImages(ArrayRef<DeviceFile> Images) {
+  DenseMap<OffloadKind, SmallVector<DeviceFile, 2>> ImagesForKind;
+  for (const DeviceFile &Image : Images)
+    ImagesForKind[Image.Kind].push_back(Image);
+
+  SmallVector<std::string, 2> WrappedImages;
+  for (const auto &KindAndImages : ImagesForKind) {
     LLVMContext Context;
     Module M("offload.wrapper.module", Context);
     M.setTargetTriple(HostTriple);
-  if (Error Err = wrapBinaries(M, ImagesToWrap))
+
+    switch (KindAndImages.getFirst()) {
+    case OFK_OpenMP:
+      if (Error Err = wrapOpenMPImages(M, KindAndImages.getSecond()))
         return std::move(Err);
+      break;
+    case OFK_Cuda:
+      if (Error Err = wrapCudaImages(M, KindAndImages.getSecond()))
+        return std::move(Err);
+      break;
+    default:
+      return createStringError(inconvertibleErrorCode(),
+                               getOffloadKindName(KindAndImages.getFirst()) +
+                                   " wrapping is not supported");
+    }
 
     if (PrintWrappedModule)
       llvm::errs() << M;
 
-  return compileModule(M);
+    auto FileOrErr = compileModule(M);
+    if (!FileOrErr)
+      return FileOrErr.takeError();
+    WrappedImages.push_back(*FileOrErr);
+  }
+
+  return WrappedImages;
 }
 
 Optional<std::string> findFile(StringRef Dir, const Twine &Name) {
   SmallString<128> Path;
   if (Dir.startswith("="))
     sys::path::append(Path, Sysroot, Dir.substr(1), Name);
   else
     sys::path::append(Path, Dir, Name);
for (std::string &Arg : LinkerArgs) {
     }
   }
 
   // Add the device bitcode libraries to the device files if any were passed in.
   for (StringRef LibraryStr : BitcodeLibraries)
     DeviceFiles.push_back(getBitcodeLibrary(LibraryStr));
 
   // Link the device images extracted from the linker input.
-  SmallVector<std::string, 16> LinkedImages;
+  SmallVector<DeviceFile, 4> LinkedImages;
   if (Error Err = linkDeviceFiles(DeviceFiles, LinkedImages))
     return reportError(std::move(Err));
 
   // Wrap each linked device image into a linkable host binary and add it to the
   // link job's inputs.
   auto FileOrErr = wrapDeviceImages(LinkedImages);
   if (!FileOrErr)
     return reportError(FileOrErr.takeError());
-  LinkerArgs.push_back(*FileOrErr);
+  LinkerArgs.append(*FileOrErr);
 
   // Run the host linking job.
   if (Error Err = runLinker(LinkerUserPath, LinkerArgs))
     return reportError(std::move(Err));
 
   // Remove the temporary files created.
   for (const auto &TempFile : TempFiles)
     if (std::error_code EC = sys::fs::remove(TempFile))
       reportError(createFileError(TempFile, EC));
 
   return EXIT_SUCCESS;
 }

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

-//===- OffloadWrapper.h -------------------------------------------*- C++ -*-===//
+//===- 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"
 
 /// Wrap the input device images into the module \p M as global symbols and
 /// registers the images with the OpenMP Offloading runtime libomptarget.
-llvm::Error wrapBinaries(llvm::Module &M,
+llvm::Error wrapOpenMPBinaries(llvm::Module &M,
                                llvm::ArrayRef<llvm::ArrayRef<char>> Images);
 
+/// Wrap 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

void createUnregisterFunction(Module &M, GlobalVariable *BinDesc) {
 
   // Add this function to global destructors.
   // Match priority of __tgt_register_lib
   appendToGlobalDtors(M, Func, /*Priority*/ 1);
 }
 
 } // namespace
 
-Error wrapBinaries(Module &M, ArrayRef<ArrayRef<char>> Images) {
+Error wrapOpenMPBinaries(Module &M, ArrayRef<ArrayRef<char>> Images) {
   GlobalVariable *Desc = createBinDesc(M, Images);
   if (!Desc)
     return createStringError(inconvertibleErrorCode(),
                              "No binary descriptors created.");
   createRegisterFunction(M, Desc);
   createUnregisterFunction(M, Desc);
   return Error::success();
 }
+
+llvm::Error wrapCudaBinary(llvm::Module &M, llvm::ArrayRef<char> Images) {
+  return createStringError(inconvertibleErrorCode(),
+                           "Cuda wrapping is not yet supported.");
+}