[StreamExecutor] Add kernel types

Summary: Add StreamExecutor kernel types.

Reviewers: jlebar, tra

Subscribers: parallel_libs-commits

Differential Revision: https://reviews.llvm.org/D23138

llvm-svn: 277827

4 files changed, 149 insertions, 0 deletions

diff --git a/parallel-libs/streamexecutor/lib/CMakeLists.txt b/parallel-libs/streamexecutor/lib/CMakeLists.txt
index f8f0c163988..6dcb5d87edd 100644
--- a/parallel-libs/streamexecutor/lib/CMakeLists.txt
+++ b/parallel-libs/streamexecutor/lib/CMakeLists.txt
@@ -6,7 +6,9 @@ add_library(
 add_library(
     streamexecutor
     $<TARGET_OBJECTS:utils>
+    Kernel.cpp
     KernelSpec.cpp)
+target_link_libraries(streamexecutor ${llvm_libs})
 
 if(STREAM_EXECUTOR_UNIT_TESTS)
     add_subdirectory(unittests)
diff --git a/parallel-libs/streamexecutor/lib/Kernel.cpp b/parallel-libs/streamexecutor/lib/Kernel.cpp
new file mode 100644
index 00000000000..af95bbe820d
--- /dev/null
+++ b/parallel-libs/streamexecutor/lib/Kernel.cpp
@@ -0,0 +1,45 @@
+//===-- Kernel.cpp - General kernel 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 types.
+///
+//===----------------------------------------------------------------------===//
+
+#include "streamexecutor/Kernel.h"
+#include "streamexecutor/Interfaces.h"
+#include "streamexecutor/StreamExecutor.h"
+
+#include "llvm/DebugInfo/Symbolize/Symbolize.h"
+
+namespace streamexecutor {
+
+KernelBase::KernelBase(StreamExecutor *ParentExecutor, const std::string &Name,
+                       const std::string &DemangledName,
+                       std::unique_ptr<KernelInterface> Implementation)
+    : ParentExecutor(ParentExecutor), Name(Name), DemangledName(DemangledName),
+      Implementation(std::move(Implementation)) {}
+
+KernelBase::~KernelBase() = default;
+
+Expected<KernelBase> KernelBase::create(StreamExecutor *ParentExecutor,
+                                        const MultiKernelLoaderSpec &Spec) {
+  auto MaybeImplementation = ParentExecutor->getKernelImplementation(Spec);
+  if (!MaybeImplementation) {
+    return MaybeImplementation.takeError();
+  }
+  std::string Name = Spec.getKernelName();
+  std::string DemangledName =
+      llvm::symbolize::LLVMSymbolizer::DemangleName(Name, nullptr);
+  KernelBase Instance(ParentExecutor, Name, DemangledName,
+                      std::move(*MaybeImplementation));
+  return std::move(Instance);
+}
+
+} // namespace streamexecutor
diff --git a/parallel-libs/streamexecutor/lib/unittests/CMakeLists.txt b/parallel-libs/streamexecutor/lib/unittests/CMakeLists.txt
index 55717aa33c7..d88332a2ee7 100644
--- a/parallel-libs/streamexecutor/lib/unittests/CMakeLists.txt
+++ b/parallel-libs/streamexecutor/lib/unittests/CMakeLists.txt
@@ -1,4 +1,14 @@
 add_executable(
+    kernel_test
+    KernelTest.cpp)
+target_link_libraries(
+    kernel_test
+    streamexecutor
+    ${GTEST_BOTH_LIBRARIES}
+    ${CMAKE_THREAD_LIBS_INIT})
+add_test(KernelTest kernel_test)
+
+add_executable(
     kernel_spec_test
     KernelSpecTest.cpp)
 target_link_libraries(
diff --git a/parallel-libs/streamexecutor/lib/unittests/KernelTest.cpp b/parallel-libs/streamexecutor/lib/unittests/KernelTest.cpp
new file mode 100644
index 00000000000..addcb83ec64
--- /dev/null
+++ b/parallel-libs/streamexecutor/lib/unittests/KernelTest.cpp
@@ -0,0 +1,92 @@
+//===-- KernelTest.cpp - Tests for Kernel objects -------------------------===//
+//
+//                     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 Kernel.
+///
+//===----------------------------------------------------------------------===//
+
+#include <cassert>
+
+#include "streamexecutor/Interfaces.h"
+#include "streamexecutor/Kernel.h"
+#include "streamexecutor/KernelSpec.h"
+#include "streamexecutor/StreamExecutor.h"
+
+#include "llvm/ADT/STLExtras.h"
+
+#include "gtest/gtest.h"
+
+namespace {
+
+namespace se = ::streamexecutor;
+
+// A StreamExecutor that returns a dummy KernelInterface.
+//
+// During construction it creates a unique_ptr to a dummy KernelInterface and it
+// also stores a separate copy of the raw pointer that is stored by that
+// unique_ptr.
+//
+// The expectation is that the code being tested will call the
+// getKernelImplementation method and will thereby take ownership of the
+// unique_ptr, but the copy of the raw pointer will stay behind in this mock
+// object. The raw pointer copy can then be used to identify the unique_ptr in
+// its new location (by comparing the raw pointer with unique_ptr::get), to
+// verify that the unique_ptr ended up where it was supposed to be.
+class MockStreamExecutor : public se::StreamExecutor {
+public:
+  MockStreamExecutor()
+      : Unique(llvm::make_unique<se::KernelInterface>()), Raw(Unique.get()) {}
+
+  // Moves the unique pointer into the returned se::Expected instance.
+  //
+  // Asserts that it is not called again after the unique pointer has been moved
+  // out.
+  se::Expected<std::unique_ptr<se::KernelInterface>>
+  getKernelImplementation(const se::MultiKernelLoaderSpec &) override {
+    assert(Unique && "MockStreamExecutor getKernelImplementation should not be "
+                     "called more than once");
+    return std::move(Unique);
+  }
+
+  // Gets the copy of the raw pointer from the original unique pointer.
+  const se::KernelInterface *getRaw() const { return Raw; }
+
+private:
+  std::unique_ptr<se::KernelInterface> Unique;
+  const se::KernelInterface *Raw;
+};
+
+// Test fixture class for typed tests for KernelBase.getImplementation.
+//
+// The only purpose of this class is to provide a name that types can be bound
+// to in the gtest infrastructure.
+template <typename T> class GetImplementationTest : public ::testing::Test {};
+
+// Types used with the GetImplementationTest fixture class.
+typedef ::testing::Types<se::KernelBase, se::TypedKernel<>,
+                         se::TypedKernel<int>>
+    GetImplementationTypes;
+
+TYPED_TEST_CASE(GetImplementationTest, GetImplementationTypes);
+
+// Tests that the kernel create functions properly fetch the implementation
+// pointers for the kernel objects they construct from the passed-in
+// StreamExecutor objects.
+TYPED_TEST(GetImplementationTest, SetImplementationDuringCreate) {
+  se::MultiKernelLoaderSpec Spec;
+  MockStreamExecutor MockExecutor;
+
+  auto MaybeKernel = TypeParam::create(&MockExecutor, Spec);
+  EXPECT_TRUE(static_cast<bool>(MaybeKernel));
+  se::KernelInterface *Implementation = MaybeKernel->getImplementation();
+  EXPECT_EQ(MockExecutor.getRaw(), Implementation);
+}
+
+} // namespace