[StreamExecutor] Add KernelLoaderSpec

[jhen]

Summary

Add definitions for the KernelLoaderSpec and MultiKernelLoaderSpec
classes to StreamExecutor. Instances of these classes are generated by the
compiler in order to provide host code with a handle to device code.

[jlebar]

streamexecutor/include/streamexecutor/KernelSpec.h

11 ↗

(On Diff #66390)

Nit: "A Foo object" or "Foo is a class".

12 ↗

(On Diff #66390)

suggest s/single/particular/

15 ↗

(On Diff #66390)

Oh, hm, I see why you said "A Foo class". I think maybe you want to say "FileLoaderSpec is the base class for types that know where to find the code."

28 ↗

(On Diff #66390)

I think this would be much easier to read if it were in second person. As-is the subject of the first phrase is MultiKernelLoaderSpec, which isn't right.

42 ↗

(On Diff #66390)

Would it be helpful to see how this function is implemented?

50 ↗

(On Diff #66390)

UserKernelLoaderSpec()?

51 ↗

(On Diff #66390)

This won't compile because MultiKernelLoaderSpec doesn't have this function.

91 ↗

(On Diff #66390)

I'm not sure we need this comment; it's pretty clear.

98 ↗

(On Diff #66390)

operator= usually returns the KernelLoaderSpec&. I think it will work as you've written it here [1], but maybe changing it will let others avoid having to look it up.

[1] http://en.cppreference.com/w/cpp/language/copy_assignment

101 ↗

(On Diff #66390)

Maybe "resides in memory as a null-terminated string."

102 ↗

(On Diff #66390)

Should this have "Spec" in the name?

106 ↗

(On Diff #66390)

A tuple as opposed to std::pair? pair also defines operator<, and seems simpler.

121 ↗

(On Diff #66390)

I don't really understand this -- we're only giving one PTX string, so of course it's going to be the default?

125 ↗

(On Diff #66390)

Since PTX source code usually (always?) has a directive at the beginning specifying the compute capability, it doesn't seem unreasonable for someone to assume that if they call this constructor, we will infer the CC from the given string.

It also seems weird that this constructor is the only way to specify a minimum-CC PTX string.

I wonder if we we need this constructor at all? If we had a public way to get the minimum CC, then we could just have the one constructor?

134 ↗

(On Diff #66390)

I wonder if it would make more sense to put the name first.

134 ↗

(On Diff #66390)

Using an initializer_list as opposed to, say, a vector, seems kind of limiting, since you can't dynamically construct an initializer_list.

(llvm::ArrayRef is really what we want, but I don't know if you want to force the compiler to pull that in.)

149 ↗

(On Diff #66390)

Since PTX is backwards-compatible, we want to return the PTX with the largest version <= the given compute-capability, right? Not clear from the comment if that's what we're doing.

154 ↗

(On Diff #66390)

Can we write this as (cc_major, cc_minor) and (2, 0), etc?

157 ↗

(On Diff #66390)

Not sure this comment is necessary here. It's kind of obvious, but if necessary maybe we can place it next to the code that uses this assumption?

177 ↗

(On Diff #66390)

Do you think it would make sense to take the size of the cubin data? I guess we don't ever need it so it doesn't matter.

Maybe I was mislead by the comments saying that we don't make a copy -- clearly we can't if we don't even know the length of the thing!

187 ↗

(On Diff #66390)

It seems kind of inconsistent that the PTX spec encompasses multiple PTX strings compiled for different archs, while the cubin one doesn't.

Unless you mean s/cubin/fatbin/? AIUI a cubin encompasses SASS code for a particular sm, while a fatbin is a combination of one or more cubins and ptx strings.

189 ↗

(On Diff #66390)

If this is null-terminated, maybe we should say that?

204 ↗

(On Diff #66390)

It looks like we do make a copy? That seems inconsistent though.

222 ↗

(On Diff #66390)

It seems like we could at least assert() that they have the same name?

263 ↗

(On Diff #66390)

extern "C" doesn't need a scope:

extern "C" void foo()

so maybe this is better said as 'if it is not declared extern "C".'

283 ↗

(On Diff #66390)

Same comments as on the constructors for those types.

293 ↗

(On Diff #66390)

Not sure these comments are necessary; seems obvious.

streamexecutor/lib/KernelSpec.cpp

27 ↗

(On Diff #66390)

Seems obvious from the code?

71 ↗

(On Diff #66390)

For chaining purposes it's more common to return *this (see e.g. operator=).

[jhen]

• Responding to comments from jlebar

streamexecutor/include/streamexecutor/KernelSpec.h

28 ↗

(On Diff #66390)

I didn't switch to second person, but I do think I cleaned up the language to make it much more understandable.

42 ↗

(On Diff #66390)

Oops, I shouldn't have put the parens here. This is supposed to be instantiating a MultiKernelLoaderSpec called UserKernelLoaderSpec.

50 ↗

(On Diff #66390)

I think the confusion here is the same as the point above. UserKernelLoaderSpec is a class instance, not a function to create a class instance.

51 ↗

(On Diff #66390)

I'm not sure what's causing the confusion here. It looks to me like I did declare MultiKernelLoaderSpec::addCUDACubinInMemory. Is that the function you are talking about?

121 ↗

(On Diff #66390)

Yes, this whole method doesn't really make sense. I removed it. I also removed the method for returning the "default" code. We can add that in later if we have a use case, but it's not needed for now.

125 ↗

(On Diff #66390)

We could add code to make this interface safer at the expense of more complexity, but I think the original design is the right trade-off because this is only meant to be called by the compiler.

149 ↗

(On Diff #66390)

I added to the comment to explain that there is no fallback logic in this function. If the exact compute capability is not found, it returns an error value. If we need that other functionality, we can add in another method later, but this is the only functionality that is present right now.

157 ↗

(On Diff #66390)

I removed the comment because the code that used it is gone now.

177 ↗

(On Diff #66390)

I don't think it would be good to take the size because we would never use it. I'm not sure what's best to put in the comment to reduce confusion. If the user thought it was a null-terminated string (which it isn't, but they might think it was), then they might be imagining that a copy is made even though no length is provided.

187 ↗

(On Diff #66390)

Good catch! What this code has been calling "cubin" is really "fatbin". That's how it is treated when it is loaded. So I changed the names to fatbin. Now they are accurate.

204 ↗

(On Diff #66390)

Oops, thanks, I missed that. I removed the copy to make it consistent with the others.

[jhen]

In addition to responding to jlebar's comments, I also added unit tests.

[jlebar]

streamexecutor/include/streamexecutor/KernelSpec.h

29 ↗

(On Diff #66559)

lgtm

52 ↗

(On Diff #66559)

I have no idea what that comment was about, sorry. I think I was just very confused by that most vexing parse.

106 ↗

(On Diff #66559)

Not sure we need this line anymore.

120 ↗

(On Diff #66559)

Here and elsewhere, I think we should change "holds a reference to". For example, if I have a std::shared_ptr<Foo>, I "hold a (strong) reference to Foo", which is the opposite of what you mean.

Maybe just "Does not take ownership of the PTXCode pointers in SpecList's elements." or something would make it clear what we're going after.

126 ↗

(On Diff #66559)

I think the original design is the right trade-off because this is only meant to be called by the compiler.

sgtm

144 ↗

(On Diff #66559)

It's not a null-terminated string, I think. I thought it was at first, but I think I was being mislead by "char" -- it's more likely that the first 4 bytes are the length (or something).

178 ↗

(On Diff #66559)

Hm...what if we used char* for text strings and uint8_t* for byte data?

streamexecutor/lib/unittests/KernelSpecTest.cpp

26 ↗

(On Diff #66559)

ASSERT has an implicit "return", so we usually call EXPECT unless the test is going to crash if the condition fails.

74 ↗

(On Diff #66559)

EXPECT_DEBUG_DEATH and you don't need the ifdefs. Same below.

[jhen]

• Respond to jlebar's comments, round 2

streamexecutor/include/streamexecutor/KernelSpec.h

144 ↗

(On Diff #66559)

Oops. I just wrote that comment in the wrong place. You're right that the char types are making things confusing.

178 ↗

(On Diff #66559)

Yes, that is much better. How do you feel about void* rather than uint8_t? I think that would be good because the data is a formless blob.

streamexecutor/lib/unittests/KernelSpecTest.cpp

74 ↗

(On Diff #66559)

Ooooooo, that's nice. Thanks!

[jlebar]

All right!

streamexecutor/include/streamexecutor/KernelSpec.h

179 ↗

(On Diff #66596)

Even better.

parallel-libs/trunk/streamexecutor/CMakeLists.txt

if(STREAM_EXECUTOR_STANDALONE)
     # Find the libraries that correspond to the LLVM components
     # that we wish to use
     llvm_map_components_to_libnames(llvm_libs support)
 
     if(STREAM_EXECUTOR_UNIT_TESTS)
         enable_testing()
         find_package(GTest REQUIRED)
         include_directories(${GTEST_INCLUDE_DIRS})
+        find_package(Threads REQUIRED)
     endif()
 else(NOT STREAM_EXECUTOR_STANDALONE)
     if(STREAM_EXECUTOR_UNIT_TESTS)
         include_directories(
             "${LLVM_MAIN_SRC_DIR}/utils/unittest/googletest/include")
     endif()
 endif(STREAM_EXECUTOR_STANDALONE)
 
 # Insist on C++ 11 features.
 set(CMAKE_CXX_STANDARD 11)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 
 # Add warning flags.
 set(CMAKE_CXX_FLAGS  "${CMAKE_CXX_FLAGS} -Wall -Wextra -Wno-unused-parameter")
 
 add_subdirectory(lib)

parallel-libs/trunk/streamexecutor/include/streamexecutor/KernelSpec.h

+//===-- KernelSpec.h - Kernel loader spec types -----------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// KernelLoaderSpec is the base class for types that know where to find the
+/// code for a data-parallel kernel in a particular format on a particular
+/// platform. So, for example, there will be one subclass that deals with CUDA
+/// PTX code, another subclass that deals with CUDA fatbin code, and yet another
+/// subclass that deals with OpenCL text code.
+///
+/// A MultiKernelLoaderSpec is basically a collection of KernelLoaderSpec
+/// instances. This is useful when code is available for the same kernel in
+/// several different formats or targeted for several different platforms. All
+/// the various KernelLoaderSpec instances for this kernel can be combined
+/// together in one MultiKernelLoaderSpec and the specific platform consumer can
+/// decide which instance of the code it wants to use.
+///
+/// MultiKernelLoaderSpec provides several helper functions to build and
+/// register KernelLoaderSpec instances all in a single operation. For example,
+/// MultiKernelLoaderSpec::addCUDAPTXInMemory can be used to construct and
+/// register a CUDAPTXInMemorySpec KernelLoaderSpec.
+///
+/// The loader spec classes declared here are designed primarily to be
+/// instantiated by the compiler, but they can also be instantiated directly by
+/// the user. A simplified example workflow which a compiler might follow in the
+/// case of a CUDA kernel that is compiled to CUDA fatbin code is as follows:
+///
+/// 1. The user defines a kernel function called UserKernel.
+/// 2. The compiler compiles the kernel code into CUDA fatbin data and embeds
+///    that data into the host code at address __UserKernelFatbinAddress.
+/// 3. The compiler adds code at the beginning of the host code to instantiate a
+///    MultiKernelLoaderSpec:
+///    \code
+///    namespace compiler_cuda_namespace {
+///      MultiKernelLoaderSpec UserKernelLoaderSpec;
+///    } // namespace compiler_cuda_namespace
+///    \endcode
+/// 4. The compiler then adds code to the host code to add the fatbin data to
+///    the new MultiKernelLoaderSpec, and to associate that data with the kernel
+///    name "UserKernel":
+///    \code
+///    namespace compiler_cuda_namespace {
+///      UserKernelLoaderSpec.addCUDAFatbinInMemory(
+///        __UserKernelFatbinAddress, "UserKernel");
+///    } // namespace compiler_cuda_namespace
+///    \encode
+/// 5. The host code, having known beforehand that the compiler would initialize
+///    a MultiKernelLoaderSpec based on the name of the CUDA kernel, makes use
+///    of the symbol cudanamespace::UserKernelLoaderSpec without defining it.
+///
+/// In the example above, the MultiKernelLoaderSpec instance created by the
+/// compiler can be used by the host code to create StreamExecutor kernel
+/// objects. In turn, those StreamExecutor kernel objects can be used by the
+/// host code to launch the kernel on the device as desired.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef STREAMEXECUTOR_KERNELSPEC_H
+#define STREAMEXECUTOR_KERNELSPEC_H
+
+#include <cassert>
+#include <map>
+#include <memory>
+#include <string>
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/StringRef.h"
+
+namespace streamexecutor {
+
+/// An object that knows how to find the code for a device kernel.
+///
+/// This is the base class for the hierarchy of loader specs. The different
+/// subclasses know how to find code in different formats (e.g. CUDA PTX, OpenCL
+/// binary).
+///
+/// This base class has functionality for storing and getting the name of the
+/// kernel as a string.
+class KernelLoaderSpec {
+public:
+  /// Returns the name of the kernel this spec loads.
+  const std::string &getKernelName() const { return KernelName; }
+
+protected:
+  explicit KernelLoaderSpec(llvm::StringRef KernelName);
+
+private:
+  std::string KernelName;
+
+  KernelLoaderSpec(const KernelLoaderSpec &) = delete;
+  KernelLoaderSpec &operator=(const KernelLoaderSpec &) = delete;
+};
+
+/// A KernelLoaderSpec for CUDA PTX code that resides in memory as a
+/// null-terminated string.
+class CUDAPTXInMemorySpec : public KernelLoaderSpec {
+public:
+  /// First component is major version, second component is minor version.
+  using ComputeCapability = std::pair<int, int>;
+
+  /// PTX code combined with its compute capability.
+  struct PTXSpec {
+    ComputeCapability TheComputeCapability;
+    const char *PTXCode;
+  };
+
+  /// Creates a CUDAPTXInMemorySpec from an array of PTXSpec objects.
+  ///
+  /// Adds each item in SpecList to this object.
+  ///
+  /// Does not take ownership of the PTXCode pointers in the SpecList elements.
+  CUDAPTXInMemorySpec(llvm::StringRef KernelName,
+                      const llvm::ArrayRef<PTXSpec> SpecList);
+
+  /// Returns a pointer to the PTX code for the requested compute capability.
+  ///
+  /// Returns nullptr on failed lookup (if the requested compute capability is
+  /// not available). Matches exactly the specified compute capability. Doesn't
+  /// try to do anything smart like finding the next best compute capability if
+  /// the specified capability cannot be found.
+  const char *getCode(int ComputeCapabilityMajor,
+                      int ComputeCapabilityMinor) const;
+
+private:
+  /// PTX code contents in memory.
+  ///
+  /// The key is a pair (cc_major, cc_minor), i.e., (2, 0), (3, 0), (3, 5).
+  std::map<ComputeCapability, const char *> PTXByComputeCapability;
+
+  CUDAPTXInMemorySpec(const CUDAPTXInMemorySpec &) = delete;
+  CUDAPTXInMemorySpec &operator=(const CUDAPTXInMemorySpec &) = delete;
+};
+
+/// A KernelLoaderSpec for CUDA fatbin code that resides in memory.
+class CUDAFatbinInMemorySpec : public KernelLoaderSpec {
+public:
+  /// Creates a CUDAFatbinInMemorySpec with a reference to the given fatbin
+  /// bytes.
+  ///
+  /// Does not take ownership of the Bytes pointer.
+  CUDAFatbinInMemorySpec(llvm::StringRef KernelName, const void *Bytes);
+
+  /// Gets the fatbin data bytes.
+  const void *getBytes() const { return Bytes; }
+
+private:
+  const void *Bytes;
+
+  CUDAFatbinInMemorySpec(const CUDAFatbinInMemorySpec &) = delete;
+  CUDAFatbinInMemorySpec &operator=(const CUDAFatbinInMemorySpec &) = delete;
+};
+
+/// A KernelLoaderSpec for OpenCL text that resides in memory as a
+/// null-terminated string.
+class OpenCLTextInMemorySpec : public KernelLoaderSpec {
+public:
+  /// Creates a OpenCLTextInMemorySpec with a reference to the given OpenCL text
+  /// code bytes.
+  ///
+  /// Does not take ownership of the Text pointer.
+  OpenCLTextInMemorySpec(llvm::StringRef KernelName, const char *Text);
+
+  /// Returns the OpenCL text contents.
+  const char *getText() const { return Text; }
+
+private:
+  const char *Text;
+
+  OpenCLTextInMemorySpec(const OpenCLTextInMemorySpec &) = delete;
+  OpenCLTextInMemorySpec &operator=(const OpenCLTextInMemorySpec &) = delete;
+};
+
+/// An object to store several different KernelLoaderSpecs for the same kernel.
+///
+/// This allows code in different formats and for different platforms to be
+/// stored all together for a single kernel.
+///
+/// Various methods are available to add a new KernelLoaderSpec to a
+/// MultiKernelLoaderSpec. There are also methods to query which formats and
+/// platforms are supported by the currently added KernelLoaderSpec objects, and
+/// methods to get the KernelLoaderSpec objects for each format and platform.
+///
+/// Since all stored KernelLoaderSpecs are supposed to reference the same
+/// kernel, they are all assumed to take the same number and type of parameters,
+/// but no checking is done to enforce this. In debug mode, all
+/// KernelLoaderSpecs are checked to make sure they have the same kernel name,
+/// so passing in specs with different kernel names can cause the program to
+/// abort.
+///
+/// This interface is prone to errors, so it is better to leave
+/// MultiKernelLoaderSpec creation and initialization to the compiler rather
+/// than doing it by hand.
+class MultiKernelLoaderSpec {
+public:
+  // Convenience getters for testing whether these platform variants have
+  // kernel loader specifications available.
+
+  bool hasCUDAPTXInMemory() const { return TheCUDAPTXInMemorySpec != nullptr; }
+  bool hasCUDAFatbinInMemory() const {
+    return TheCUDAFatbinInMemorySpec != nullptr;
+  }
+  bool hasOpenCLTextInMemory() const {
+    return TheOpenCLTextInMemorySpec != nullptr;
+  }
+
+  // Accessors for platform variant kernel load specifications.
+  //
+  // Precondition: corresponding has* method returns true.
+
+  const CUDAPTXInMemorySpec &getCUDAPTXInMemory() const {
+    assert(hasCUDAPTXInMemory() && "getting spec that is not present");
+    return *TheCUDAPTXInMemorySpec;
+  }
+  const CUDAFatbinInMemorySpec &getCUDAFatbinInMemory() const {
+    assert(hasCUDAFatbinInMemory() && "getting spec that is not present");
+    return *TheCUDAFatbinInMemorySpec;
+  }
+  const OpenCLTextInMemorySpec &getOpenCLTextInMemory() const {
+    assert(hasOpenCLTextInMemory() && "getting spec that is not present");
+    return *TheOpenCLTextInMemorySpec;
+  }
+
+  // Builder-pattern-like methods for use in initializing a
+  // MultiKernelLoaderSpec.
+  //
+  // Each of these should be used at most once for a single
+  // MultiKernelLoaderSpec object. See file comment for example usage.
+  //
+  // Note that the KernelName parameter must be consistent with the kernel in
+  // the PTX or OpenCL being loaded. Also be aware that in CUDA C++ the kernel
+  // name may be mangled by the compiler if it is not declared extern "C".
+
+  /// Does not take ownership of the PTXCode pointers in the SpecList elements.
+  MultiKernelLoaderSpec &
+  addCUDAPTXInMemory(llvm::StringRef KernelName,
+                     llvm::ArrayRef<CUDAPTXInMemorySpec::PTXSpec> SpecList);
+
+  /// Does not take ownership of the FatbinBytes pointer.
+  MultiKernelLoaderSpec &addCUDAFatbinInMemory(llvm::StringRef KernelName,
+                                               const void *FatbinBytes);
+
+  /// Does not take ownership of the OpenCLText pointer.
+  MultiKernelLoaderSpec &addOpenCLTextInMemory(llvm::StringRef KernelName,
+                                               const char *OpenCLText);
+
+private:
+  void setKernelName(llvm::StringRef KernelName);
+
+  std::unique_ptr<std::string> TheKernelName;
+  std::unique_ptr<CUDAPTXInMemorySpec> TheCUDAPTXInMemorySpec;
+  std::unique_ptr<CUDAFatbinInMemorySpec> TheCUDAFatbinInMemorySpec;
+  std::unique_ptr<OpenCLTextInMemorySpec> TheOpenCLTextInMemorySpec;
+};
+
+} // namespace streamexecutor
+
+#endif // STREAMEXECUTOR_KERNELSPEC_H

parallel-libs/trunk/streamexecutor/lib/CMakeLists.txt

 add_library(
     utils
     OBJECT
     Utils/Error.cpp)
+
+add_library(
+    streamexecutor
+    $<TARGET_OBJECTS:utils>
+    KernelSpec.cpp)
+
+if(STREAM_EXECUTOR_UNIT_TESTS)
+    add_subdirectory(unittests)
+endif()

parallel-libs/trunk/streamexecutor/lib/KernelSpec.cpp

+//===-- KernelSpec.cpp - General kernel spec implementation ---------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// This file contains the implementation details for kernel loader specs.
+///
+//===----------------------------------------------------------------------===//
+
+#include "streamexecutor/KernelSpec.h"
+
+#include "llvm/ADT/STLExtras.h"
+
+namespace streamexecutor {
+
+KernelLoaderSpec::KernelLoaderSpec(llvm::StringRef KernelName)
+    : KernelName(KernelName) {}
+
+CUDAPTXInMemorySpec::CUDAPTXInMemorySpec(
+    llvm::StringRef KernelName,
+    const llvm::ArrayRef<CUDAPTXInMemorySpec::PTXSpec> SpecList)
+    : KernelLoaderSpec(KernelName) {
+  for (const auto &Spec : SpecList) {
+    PTXByComputeCapability.emplace(Spec.TheComputeCapability, Spec.PTXCode);
+  }
+}
+
+const char *CUDAPTXInMemorySpec::getCode(int ComputeCapabilityMajor,
+                                         int ComputeCapabilityMinor) const {
+  auto PTXIter =
+      PTXByComputeCapability.find(CUDAPTXInMemorySpec::ComputeCapability{
+          ComputeCapabilityMajor, ComputeCapabilityMinor});
+  if (PTXIter == PTXByComputeCapability.end()) {
+    return nullptr;
+  }
+  return PTXIter->second;
+}
+
+CUDAFatbinInMemorySpec::CUDAFatbinInMemorySpec(llvm::StringRef KernelName,
+                                               const void *Bytes)
+    : KernelLoaderSpec(KernelName), Bytes(Bytes) {}
+
+OpenCLTextInMemorySpec::OpenCLTextInMemorySpec(llvm::StringRef KernelName,
+                                               const char *Text)
+    : KernelLoaderSpec(KernelName), Text(Text) {}
+
+void MultiKernelLoaderSpec::setKernelName(llvm::StringRef KernelName) {
+  if (TheKernelName) {
+    assert(KernelName.equals(*TheKernelName) &&
+           "different kernel names in one MultiKernelLoaderSpec");
+  } else {
+    TheKernelName = llvm::make_unique<std::string>(KernelName);
+  }
+}
+
+MultiKernelLoaderSpec &MultiKernelLoaderSpec::addCUDAPTXInMemory(
+    llvm::StringRef KernelName,
+    llvm::ArrayRef<CUDAPTXInMemorySpec::PTXSpec> SpecList) {
+  assert((TheCUDAPTXInMemorySpec == nullptr) &&
+         "illegal loader spec overwrite");
+  setKernelName(KernelName);
+  TheCUDAPTXInMemorySpec =
+      llvm::make_unique<CUDAPTXInMemorySpec>(KernelName, SpecList);
+  return *this;
+}
+
+MultiKernelLoaderSpec &
+MultiKernelLoaderSpec::addCUDAFatbinInMemory(llvm::StringRef KernelName,
+                                             const void *Bytes) {
+  assert((TheCUDAFatbinInMemorySpec == nullptr) &&
+         "illegal loader spec overwrite");
+  setKernelName(KernelName);
+  TheCUDAFatbinInMemorySpec =
+      llvm::make_unique<CUDAFatbinInMemorySpec>(KernelName, Bytes);
+  return *this;
+}
+
+MultiKernelLoaderSpec &
+MultiKernelLoaderSpec::addOpenCLTextInMemory(llvm::StringRef KernelName,
+                                             const char *OpenCLText) {
+  assert((TheOpenCLTextInMemorySpec == nullptr) &&
+         "illegal loader spec overwrite");
+  setKernelName(KernelName);
+  TheOpenCLTextInMemorySpec =
+      llvm::make_unique<OpenCLTextInMemorySpec>(KernelName, OpenCLText);
+  return *this;
+}
+
+} // namespace streamexecutor

parallel-libs/trunk/streamexecutor/lib/unittests/CMakeLists.txt

+add_executable(
+    kernel_spec_test
+    KernelSpecTest.cpp)
+target_link_libraries(
+    kernel_spec_test
+    streamexecutor
+    ${GTEST_BOTH_LIBRARIES}
+    ${CMAKE_THREAD_LIBS_INIT})
+add_test(KernelSpecTest kernel_spec_test)

parallel-libs/trunk/streamexecutor/lib/unittests/KernelSpecTest.cpp

+//===-- KernelSpecTest.cpp - Tests for KernelSpec -------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// This file contains the unit tests for the code in KernelSpec.
+///
+//===----------------------------------------------------------------------===//
+
+#include "streamexecutor/KernelSpec.h"
+
+#include "gtest/gtest.h"
+
+namespace {
+
+namespace se = ::streamexecutor;
+
+TEST(CUDAPTXInMemorySpec, NoCode) {
+  se::CUDAPTXInMemorySpec Spec("KernelName", {});
+  EXPECT_EQ("KernelName", Spec.getKernelName());
+  EXPECT_EQ(nullptr, Spec.getCode(1, 0));
+}
+
+TEST(CUDAPTXInMemorySpec, SingleComputeCapability) {
+  const char *PTXCodeString = "Dummy PTX code";
+  se::CUDAPTXInMemorySpec Spec("KernelName", {{{1, 0}, PTXCodeString}});
+  EXPECT_EQ("KernelName", Spec.getKernelName());
+  EXPECT_EQ(PTXCodeString, Spec.getCode(1, 0));
+  EXPECT_EQ(nullptr, Spec.getCode(2, 0));
+}
+
+TEST(CUDAPTXInMemorySpec, TwoComputeCapabilities) {
+  const char *PTXCodeString10 = "Dummy PTX code 10";
+  const char *PTXCodeString30 = "Dummy PTX code 30";
+  se::CUDAPTXInMemorySpec Spec(
+      "KernelName", {{{1, 0}, PTXCodeString10}, {{3, 0}, PTXCodeString30}});
+  EXPECT_EQ("KernelName", Spec.getKernelName());
+  EXPECT_EQ(PTXCodeString10, Spec.getCode(1, 0));
+  EXPECT_EQ(PTXCodeString30, Spec.getCode(3, 0));
+  EXPECT_EQ(nullptr, Spec.getCode(2, 0));
+}
+
+TEST(CUDAFatbinInMemorySpec, BasicUsage) {
+  const char *FatbinBytes = "Dummy fatbin bytes";
+  se::CUDAFatbinInMemorySpec Spec("KernelName", FatbinBytes);
+  EXPECT_EQ("KernelName", Spec.getKernelName());
+  EXPECT_EQ(FatbinBytes, Spec.getBytes());
+}
+
+TEST(OpenCLTextInMemorySpec, BasicUsage) {
+  const char *OpenCLText = "Dummy OpenCL text";
+  se::OpenCLTextInMemorySpec Spec("KernelName", OpenCLText);
+  EXPECT_EQ("KernelName", Spec.getKernelName());
+  EXPECT_EQ(OpenCLText, Spec.getText());
+}
+
+TEST(MultiKernelLoaderSpec, NoCode) {
+  se::MultiKernelLoaderSpec MultiSpec;
+  EXPECT_FALSE(MultiSpec.hasCUDAPTXInMemory());
+  EXPECT_FALSE(MultiSpec.hasCUDAFatbinInMemory());
+  EXPECT_FALSE(MultiSpec.hasOpenCLTextInMemory());
+
+  EXPECT_DEBUG_DEATH(MultiSpec.getCUDAPTXInMemory(),
+                     "getting spec that is not present");
+  EXPECT_DEBUG_DEATH(MultiSpec.getCUDAFatbinInMemory(),
+                     "getting spec that is not present");
+  EXPECT_DEBUG_DEATH(MultiSpec.getOpenCLTextInMemory(),
+                     "getting spec that is not present");
+}
+
+TEST(MultiKernelLoaderSpec, Registration) {
+  se::MultiKernelLoaderSpec MultiSpec;
+  const char *KernelName = "KernelName";
+  const char *PTXCodeString = "Dummy PTX code";
+  const char *FatbinBytes = "Dummy fatbin bytes";
+  const char *OpenCLText = "Dummy OpenCL text";
+
+  MultiSpec.addCUDAPTXInMemory(KernelName, {{{1, 0}, PTXCodeString}})
+      .addCUDAFatbinInMemory(KernelName, FatbinBytes)
+      .addOpenCLTextInMemory(KernelName, OpenCLText);
+
+  EXPECT_TRUE(MultiSpec.hasCUDAPTXInMemory());
+  EXPECT_TRUE(MultiSpec.hasCUDAFatbinInMemory());
+  EXPECT_TRUE(MultiSpec.hasOpenCLTextInMemory());
+
+  EXPECT_EQ(KernelName, MultiSpec.getCUDAPTXInMemory().getKernelName());
+  EXPECT_EQ(PTXCodeString, MultiSpec.getCUDAPTXInMemory().getCode(1, 0));
+  EXPECT_EQ(nullptr, MultiSpec.getCUDAPTXInMemory().getCode(2, 0));
+
+  EXPECT_EQ(KernelName, MultiSpec.getCUDAFatbinInMemory().getKernelName());
+  EXPECT_EQ(FatbinBytes, MultiSpec.getCUDAFatbinInMemory().getBytes());
+
+  EXPECT_EQ(KernelName, MultiSpec.getOpenCLTextInMemory().getKernelName());
+  EXPECT_EQ(OpenCLText, MultiSpec.getOpenCLTextInMemory().getText());
+}
+
+TEST(MultiKernelLoaderSpec, RegisterTwice) {
+  se::MultiKernelLoaderSpec MultiSpec;
+  const char *KernelName = "KernelName";
+  const char *FatbinBytes = "Dummy fatbin bytes";
+
+  MultiSpec.addCUDAFatbinInMemory(KernelName, FatbinBytes);
+
+  EXPECT_DEBUG_DEATH(MultiSpec.addCUDAFatbinInMemory(KernelName, FatbinBytes),
+                     "illegal loader spec overwrite");
+}
+
+TEST(MultiKernelLoaderSpec, ConflictingKernelNames) {
+  se::MultiKernelLoaderSpec MultiSpec;
+  const char *KernelNameA = "KernelName";
+  std::string KernelNameB = KernelNameA;
+  const char *PTXCodeString = "Dummy PTX code";
+  const char *FatbinBytes = "Dummy fatbin bytes";
+
+  // Check that names don't conflict if they are equivalent strings in different
+  // locations.
+  MultiSpec.addCUDAPTXInMemory(KernelNameA, {{{1, 0}, PTXCodeString}})
+      .addCUDAFatbinInMemory(KernelNameB, FatbinBytes);
+
+  const char *OtherKernelName = "OtherKernelName";
+  const char *OpenCLText = "Dummy OpenCL text";
+  EXPECT_DEBUG_DEATH(
+      MultiSpec.addOpenCLTextInMemory(OtherKernelName, OpenCLText),
+      "different kernel names in one MultiKernelLoaderSpec");
+}
+
+} // namespace