[SE] KernelSpec return best PTX

[jhen]

Summary

Before, the kernel spec would only return PTX for exactly the requested
compute capability. With this patch it will now return the PTX with the
largest compute capability that does not exceed that requested compute
capability.

[jlebar]

streamexecutor/lib/KernelSpec.cpp

39

Gosh this is subtle -- even though I've seen it before, I still had to spend five minutes convincing myself it's right. Oh well, I don't have a concrete suggestion. :)

streamexecutor/include/streamexecutor/KernelSpec.h

public:
   ///
   /// 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<CUDAPTXInMemorySpec::PTXSpec> SpecList);
 
-  /// Returns a pointer to the PTX code for the requested compute capability.
+  /// Returns a pointer to the PTX code for the greatest compute capability not
+  /// exceeding 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.
+  /// Returns nullptr on failed lookup (if the requested version is not
+  /// available and no lower versions are available).
   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;

streamexecutor/lib/KernelSpec.cpp

CUDAPTXInMemorySpec::CUDAPTXInMemorySpec(
     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{
+  auto Iterator =
+      PTXByComputeCapability.upper_bound(CUDAPTXInMemorySpec::ComputeCapability{
           ComputeCapabilityMajor, ComputeCapabilityMinor});
-  if (PTXIter == PTXByComputeCapability.end())
+  if (Iterator == PTXByComputeCapability.begin())
     return nullptr;
-  return PTXIter->second;
+  --Iterator;
+  return Iterator->second;
 }
 
 CUDAFatbinInMemorySpec::CUDAFatbinInMemorySpec(llvm::StringRef KernelName,
                                                const void *Bytes)
     : KernelLoaderSpec(KernelName), Bytes(Bytes) {}
 
 OpenCLTextInMemorySpec::OpenCLTextInMemorySpec(llvm::StringRef KernelName,
                                                const char *Text)

streamexecutor/unittests/CoreTests/KernelSpecTest.cpp

TEST(CUDAPTXInMemorySpec, NoCode) {
   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(nullptr, Spec.getCode(0, 5));
   EXPECT_EQ(PTXCodeString, Spec.getCode(1, 0));
-  EXPECT_EQ(nullptr, Spec.getCode(2, 0));
+  EXPECT_EQ(PTXCodeString, 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(nullptr, Spec.getCode(0, 5));
   EXPECT_EQ(PTXCodeString10, Spec.getCode(1, 0));
   EXPECT_EQ(PTXCodeString30, Spec.getCode(3, 0));
-  EXPECT_EQ(nullptr, Spec.getCode(2, 0));
+  EXPECT_EQ(PTXCodeString10, 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());
 }
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(nullptr, MultiSpec.getCUDAPTXInMemory().getCode(0, 5));
   EXPECT_EQ(PTXCodeString, MultiSpec.getCUDAPTXInMemory().getCode(1, 0));
-  EXPECT_EQ(nullptr, MultiSpec.getCUDAPTXInMemory().getCode(2, 0));
+  EXPECT_EQ(PTXCodeString, 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());
 }