[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

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

12

suggest s/single/particular/

15

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

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

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

50

UserKernelLoaderSpec()?

51

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

91

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

98

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

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

102

Should this have "Spec" in the name?

106

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

121

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

125

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

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

134

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

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

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

157

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

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

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

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

204

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

222

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

263

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

Same comments as on the constructors for those types.

293

Not sure these comments are necessary; seems obvious.

streamexecutor/lib/KernelSpec.cpp

27

Seems obvious from the code?

71

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

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

42

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

50

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

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

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

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

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

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

177

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

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

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

30

lgtm

53

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

107

Not sure we need this line anymore.

121

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.

127

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

sgtm

145

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).

179

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

145

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

179

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

180

Even better.

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
+/// A KernelLoaderSpec is a class that knows where to find the code for a
+/// data-parallel kernel in a single format on a single platform. So, for
+/// example, there will be one subclass that deals with CUDA PTX code, another
+/// subclass that deals with CUDA cubin 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.
+///
+/// Rather than instantiating each KernelLoaderSpec separately and then
+/// registering it with the MultiKernelLoaderSpec, the MultiKernelLoaderSpec
+/// provides several helper functions, such as
+/// MultiKernelLoaderSpec::addCUDAPTXInMemory, to construct the KernelLoaderSpec
+/// and add it to the MultiKernelLoaderSpec at the same time.
+///
+/// These loader spec classes 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 cubin code is as follows:
+///
+/// 1. The user defines a kernel function called UserKernel.
+/// 2. The compiler compiles the kernel code into cubin data and embeds that
+///    data into the host code at address __UserKernelCubinAddress.
+/// 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 cubin data to the
+///    new MultiKernelLoaderSpec, and to associate that data with the kernel
+///    name "UserKernel":
+///    \code
+///    namespace compiler_cuda_namespace {
+///      UserKernelLoaderSpec.addCUDACubinInMemory(
+///        __UserKernelCubinAddress, "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/Support/RWMutex.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:
+  /// Allows subclasses to set the kernel name during construction.
+  explicit KernelLoaderSpec(const std::string &KernelName);
+
+private:
+  std::string KernelName;
+
+  KernelLoaderSpec(const KernelLoaderSpec &) = delete;
+  void operator=(const KernelLoaderSpec &) = delete;
+};
+
+/// A KernelLoaderSpec for CUDA PTX code that resides in memory.
+class CUDAPTXInMemory : public KernelLoaderSpec {
+public:
+  /// First component is major version, second component is minor version.
+  ///
+  /// A tuple is used because this will be the key of a map.
+  using ComputeCapability = std::tuple<int, int>;
+
+  /// PTX code combined with its compute capability.
+  struct PTXSpec {
+    ComputeCapability TheComputeCapability;
+    const char *PTXCode;
+  };
+
+  /// Single-PTX constructor.
+  ///
+  /// Adds the provided PTX version with the minimum compute capability.
+  ///
+  /// Since the compute capability is unknown, the PTX is assumed to be very
+  /// generally usable - in other words, PTX specified in this manner is VERY
+  /// likely to be used as the default.
+  ///
+  /// Holds a reference to the passed in PTXCode memory. Does not make a copy
+  /// and does not take ownership.
+  CUDAPTXInMemory(const char *PTXCode, const std::string &KernelName);
+
+  /// Multiple-PTX-version constructor.
+  ///
+  /// Adds each item in spec_list to this object.
+  ///
+  /// Holds a reference to each passed in PTXCode memory. Does not make a copy
+  /// and does not take ownership.
+  CUDAPTXInMemory(const std::initializer_list<PTXSpec> &SpecList,
+                  const std::string &KernelName);
+
+  /// Returns a pointer to the PTX code with the lowest-valued compute
+  /// capability.
+  ///
+  /// For example, if PTX codes of compute capabilities 2.0, 3.0, and 3.5 are
+  /// all available, the version for compute capability 2.0 will be returned.
+  /// Returns nullptr on failed lookup (if no version is available).
+  const char *getDefaultCode() const;
+
+  /// Returns a pointer to the PTX code for the requested compute capability.
+  ///
+  /// Returns nullptr on failed lookup (if the requested version is not
+  /// available).
+  const char *getCode(int ComputeCapabilityMajor,
+                      int ComputeCapabilityMinor) const;
+
+private:
+  /// PTX code contents in memory.
+  ///
+  /// The key is a tuple "<cc_major>,<cc_minor>", i.e., "2,0", "3,0", "3,5".
+  /// Because compute capabilities represented in this way have a clear sorting
+  /// order, map::begin() will give the lowest-numbered version available, i.e.
+  /// the default.
+  std::map<ComputeCapability, const char *> PTXByComputeCapability;
+
+  /// Defines the minimum compute capability possible.
+  ///
+  /// Used when PTX has no compute capability specified (in the single-PTX
+  /// constructor).
+  static const ComputeCapability MinimumComputeCapability;
+
+  CUDAPTXInMemory(const CUDAPTXInMemory &) = delete;
+  void operator=(const CUDAPTXInMemory &) = delete;
+};
+
+/// A KernelLoaderSpec for CUDA cubin code that resides in memory.
+class CUDACubinInMemory : public KernelLoaderSpec {
+public:
+  /// Creates a CUDACubinInMemory with a reference to the given cubin bytes.
+  ///
+  /// Holds a reference to the passed in Bytes memory. Does not make a copy and
+  /// does not take ownership.
+  CUDACubinInMemory(const char *Bytes, const std::string &KernelName);
+
+  /// Gets the cubin data bytes.
+  const char *getBytes() const { return Bytes; }
+
+private:
+  const char *Bytes;
+
+  CUDACubinInMemory(const CUDACubinInMemory &) = delete;
+  void operator=(const CUDACubinInMemory &) = delete;
+};
+
+/// A KernelLoaderSpec for OpenCL text that resides in memory.
+class OpenCLTextInMemory : public KernelLoaderSpec {
+public:
+  /// Creates a OpenCLTextInMemory with a reference to the given OpenCL text
+  /// code bytes.
+  ///
+  /// Holds a reference to the passed in Text memory. Does not make a copy and
+  /// does not take ownership.
+  OpenCLTextInMemory(const char *Text, const std::string &KernelName);
+
+  /// Returns the OpenCL text contents.
+  const std::string &getText() const { return Text; }
+
+private:
+  /// OpenCL translation unit text contents in memory.
+  std::string Text;
+
+  OpenCLTextInMemory(const OpenCLTextInMemory &) = delete;
+  void operator=(const OpenCLTextInMemory &) = 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,
+/// and to have the same name, but no checking is done to enforce this. Since
+/// this interface is prone to errors, 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 TheCUDAPTXInMemory != nullptr; }
+  bool hasCUDACubinInMemory() const { return TheCUDACubinInMemory != nullptr; }
+  bool hasOpenCLTextInMemory() const {
+    return TheOpenCLTextInMemory != nullptr;
+  }
+
+  // Accessors for platform variant kernel load specifications.
+  //
+  // Precondition: corresponding has* method returns true.
+
+  const CUDACubinInMemory &getCUDACubinInMemory() const {
+    assert(hasCUDACubinInMemory());
+    return *TheCUDACubinInMemory;
+  }
+  const CUDAPTXInMemory &getCUDAPTXInMemory() const {
+    assert(hasCUDAPTXInMemory());
+    return *TheCUDAPTXInMemory;
+  }
+  const OpenCLTextInMemory &getOpenCLTextInMemory() const {
+    assert(hasOpenCLTextInMemory());
+    return *TheOpenCLTextInMemory;
+  }
+
+  // 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 in an extern "C"
+  // scope.
+
+  /// Holds a reference to the PTXCode memory. Does not make a copy and does not
+  /// take ownership.
+  MultiKernelLoaderSpec *addCUDAPTXInMemory(const char *PTXCode,
+                                            const std::string &KernelName);
+  /// Holds a reference to each PTX code memory in SpecList. Does not make a
+  /// copy and does not take ownership.
+  MultiKernelLoaderSpec *
+  addCUDAPTXInMemory(std::initializer_list<CUDAPTXInMemory::PTXSpec> SpecList,
+                     const std::string &KernelName);
+
+  /// Holds a reference to the CubinBytes memory. Does not make a copy and does
+  /// not take ownership.
+  MultiKernelLoaderSpec *addCUDACubinInMemory(const char *CubinBytes,
+                                              const std::string &KernelName);
+
+  /// Holds a reference to the OpenCLText memory. Does not make a copy and does
+  /// not take ownership.
+  MultiKernelLoaderSpec *addOpenCLTextInMemory(const char *OpenCLText,
+                                               const std::string &KernelName);
+
+private:
+  /// PTX code that resides in memory.
+  std::unique_ptr<CUDAPTXInMemory> TheCUDAPTXInMemory;
+
+  /// Binary CUDA program in memory.
+  std::unique_ptr<CUDACubinInMemory> TheCUDACubinInMemory;
+
+  /// OpenCL text that resides in memory.
+  std::unique_ptr<OpenCLTextInMemory> TheOpenCLTextInMemory;
+};
+
+} // namespace streamexecutor
+
+#endif // STREAMEXECUTOR_KERNELSPEC_H

streamexecutor/lib/CMakeLists.txt

 add_library(
     utils
     OBJECT
     Utils/Error.cpp)
+
+add_library(
+    streamexecutor
+    $<TARGET_OBJECTS:utils>
+    KernelSpec.cpp)

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 <cassert>
+#include <string>
+
+#include "llvm/ADT/STLExtras.h"
+
+namespace streamexecutor {
+
+KernelLoaderSpec::KernelLoaderSpec(const std::string &KernelName)
+    : KernelName(KernelName) {}
+
+// Minimum compute capability is 1.0.
+const CUDAPTXInMemory::ComputeCapability
+    CUDAPTXInMemory::MinimumComputeCapability{1, 0};
+
+CUDAPTXInMemory::CUDAPTXInMemory(const char *PTXCode,
+                                 const std::string &KernelName)
+    : KernelLoaderSpec(KernelName) {
+  PTXByComputeCapability.emplace(MinimumComputeCapability, PTXCode);
+}
+
+CUDAPTXInMemory::CUDAPTXInMemory(
+    const std::initializer_list<CUDAPTXInMemory::PTXSpec> &SpecList,
+    const std::string &KernelName)
+    : KernelLoaderSpec(KernelName) {
+  for (const auto &Spec : SpecList) {
+    PTXByComputeCapability.emplace(Spec.TheComputeCapability, Spec.PTXCode);
+  }
+}
+
+const char *CUDAPTXInMemory::getDefaultCode() const {
+  if (PTXByComputeCapability.empty()) {
+    return nullptr;
+  }
+  return PTXByComputeCapability.begin()->second;
+}
+
+const char *CUDAPTXInMemory::getCode(int ComputeCapabilityMajor,
+                                     int ComputeCapabilityMinor) const {
+  auto PTXIter = PTXByComputeCapability.find(CUDAPTXInMemory::ComputeCapability{
+      ComputeCapabilityMajor, ComputeCapabilityMinor});
+  if (PTXIter == PTXByComputeCapability.end()) {
+    return nullptr;
+  }
+  return PTXIter->second;
+}
+
+CUDACubinInMemory::CUDACubinInMemory(const char *Bytes,
+                                     const std::string &KernelName)
+    : KernelLoaderSpec(KernelName), Bytes(Bytes) {}
+
+OpenCLTextInMemory::OpenCLTextInMemory(const char *Text,
+                                       const std::string &KernelName)
+    : KernelLoaderSpec(KernelName), Text(Text) {}
+
+MultiKernelLoaderSpec *
+MultiKernelLoaderSpec::addCUDAPTXInMemory(const char *PTXCode,
+                                          const std::string &KernelName) {
+  assert(TheCUDAPTXInMemory == nullptr);
+  TheCUDAPTXInMemory = llvm::make_unique<CUDAPTXInMemory>(PTXCode, KernelName);
+  return this;
+}
+
+MultiKernelLoaderSpec *MultiKernelLoaderSpec::addCUDAPTXInMemory(
+    std::initializer_list<CUDAPTXInMemory::PTXSpec> SpecList,
+    const std::string &KernelName) {
+  assert(TheCUDAPTXInMemory == nullptr);
+  TheCUDAPTXInMemory = llvm::make_unique<CUDAPTXInMemory>(SpecList, KernelName);
+  return this;
+}
+
+MultiKernelLoaderSpec *
+MultiKernelLoaderSpec::addCUDACubinInMemory(const char *Bytes,
+                                            const std::string &KernelName) {
+  assert(TheCUDACubinInMemory == nullptr);
+  TheCUDACubinInMemory =
+      llvm::make_unique<CUDACubinInMemory>(Bytes, KernelName);
+  return this;
+}
+
+MultiKernelLoaderSpec *
+MultiKernelLoaderSpec::addOpenCLTextInMemory(const char *OpenCLText,
+                                             const std::string &KernelName) {
+  assert(TheOpenCLTextInMemory == nullptr);
+  TheOpenCLTextInMemory =
+      llvm::make_unique<OpenCLTextInMemory>(OpenCLText, KernelName);
+  return this;
+}
+
+} // namespace streamexecutor