1 files changed, 132 insertions, 0 deletions
diff --git a/parallel-libs/streamexecutor/unittests/CoreTests/KernelSpecTest.cpp b/parallel-libs/streamexecutor/unittests/CoreTests/KernelSpecTest.cpp
new file mode 100644
index 00000000000..fc9eb549968
--- /dev/null
+++ b/parallel-libs/streamexecutor/unittests/CoreTests/KernelSpecTest.cpp
@@ -0,0 +1,132 @@
+//===-- 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