[StreamExecutor] Add kernel types

Summary: Add StreamExecutor kernel types.

Reviewers: jlebar, tra

Subscribers: parallel_libs-commits

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

9 files changed, 383 insertions, 1 deletions

diff --git a/parallel-libs/streamexecutor/CMakeLists.txt b/parallel-libs/streamexecutor/CMakeLists.txt
index 18c28d62fef..0dd788e3cff 100644
--- a/parallel-libs/streamexecutor/CMakeLists.txt
+++ b/parallel-libs/streamexecutor/CMakeLists.txt
@@ -26,7 +26,7 @@ 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)
+    llvm_map_components_to_libnames(llvm_libs support symbolize)
 
     if(STREAM_EXECUTOR_UNIT_TESTS)
         enable_testing()
diff --git a/parallel-libs/streamexecutor/include/streamexecutor/Interfaces.h b/parallel-libs/streamexecutor/include/streamexecutor/Interfaces.h
new file mode 100644
index 00000000000..830ca0499b6
--- /dev/null
+++ b/parallel-libs/streamexecutor/include/streamexecutor/Interfaces.h
@@ -0,0 +1,29 @@
+//===-- Interfaces.h - Interfaces to platform-specific impls ----*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// Interfaces to platform-specific StreamExecutor type implementations.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef STREAMEXECUTOR_INTERFACES_H
+#define STREAMEXECUTOR_INTERFACES_H
+
+namespace streamexecutor {
+
+/// Methods supported by device kernel function objects on all platforms.
+class KernelInterface {
+  // TODO(jhen): Add methods.
+};
+
+// TODO(jhen): Add other interfaces such as Stream.
+
+} // namespace streamexecutor
+
+#endif // STREAMEXECUTOR_INTERFACES_H
diff --git a/parallel-libs/streamexecutor/include/streamexecutor/Kernel.h b/parallel-libs/streamexecutor/include/streamexecutor/Kernel.h
new file mode 100644
index 00000000000..765798afed4
--- /dev/null
+++ b/parallel-libs/streamexecutor/include/streamexecutor/Kernel.h
@@ -0,0 +1,158 @@
+//===-- Kernel.h - StreamExecutor kernel types ------------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// Types to represent device kernels (code compiled to run on GPU or other
+/// accelerator).
+///
+/// The TypedKernel class is used to provide type safety to the user API's
+/// launch functions, and the KernelBase class is used like a void* function
+/// pointer to perform type-unsafe operations inside StreamExecutor.
+///
+/// With the kernel parameter types recorded in the TypedKernel template
+/// parameters, type-safe kernel launch functions can be written with signatures
+/// like the following:
+/// \code
+///     template <typename... ParameterTs>
+///     void Launch(
+///       const TypedKernel<ParameterTs...> &Kernel, ParamterTs... Arguments);
+/// \endcode
+/// and the compiler will check that the user passes in arguments with types
+/// matching the corresponding kernel parameters.
+///
+/// A problem is that a TypedKernel template specialization with the right
+/// parameter types must be passed as the first argument to the Launch function,
+/// and it's just as hard to get the types right in that template specialization
+/// as it is to get them right for the kernel arguments.
+///
+/// With this problem in mind, it is not recommended for users to specialize the
+/// TypedKernel template class themselves, but instead to let the compiler do it
+/// for them. When the compiler encounters a device kernel function, it can
+/// create a TypedKernel template specialization in the host code that has the
+/// right parameter types for that kernel and which has a type name based on the
+/// name of the kernel function.
+///
+/// For example, if a CUDA device kernel function with the following signature
+/// has been defined:
+/// \code
+///     void Saxpy(float *A, float *X, float *Y);
+/// \endcode
+/// the compiler can insert the following declaration in the host code:
+/// \code
+///     namespace compiler_cuda_namespace {
+///     using SaxpyKernel =
+///         streamexecutor::TypedKernel<float *, float *, float *>;
+///     } // namespace compiler_cuda_namespace
+/// \endcode
+/// and then the user can launch the kernel by calling the StreamExecutor launch
+/// function as follows:
+/// \code
+///     namespace ccn = compiler_cuda_namespace;
+///     // Assumes Executor is a pointer to the StreamExecutor on which to
+///     // launch the kernel.
+///     //
+///     // See KernelSpec.h for details on how the compiler can create a
+///     // MultiKernelLoaderSpec instance like SaxpyKernelLoaderSpec below.
+///     Expected<ccn::SaxpyKernel> MaybeKernel =
+///         ccn::SaxpyKernel::create(Executor, ccn::SaxpyKernelLoaderSpec);
+///     if (!MaybeKernel) { /* Handle error */ }
+///     ccn::SaxpyKernel SaxpyKernel = *MaybeKernel;
+///     Launch(SaxpyKernel, A, X, Y);
+/// \endcode
+///
+/// With the compiler's help in specializing TypedKernel for each device kernel
+/// function (and generating a MultiKernelLoaderSpec instance for each kernel),
+/// the user can safely launch the device kernel from the host and get an error
+/// message at compile time if the argument types don't match the kernel
+/// parameter types.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef STREAMEXECUTOR_KERNEL_H
+#define STREAMEXECUTOR_KERNEL_H
+
+#include "streamexecutor/KernelSpec.h"
+#include "streamexecutor/Utils/Error.h"
+
+#include <memory>
+
+namespace streamexecutor {
+
+class KernelInterface;
+class StreamExecutor;
+
+/// The base class for device kernel functions.
+///
+/// This class has no information about the types of the parameters taken by the
+/// kernel, so it is analogous to a void* pointer to a device function.
+///
+/// See the TypedKernel class below for the subclass which does have information
+/// about parameter types.
+class KernelBase {
+public:
+  KernelBase(KernelBase &&) = default;
+  KernelBase &operator=(KernelBase &&) = default;
+  ~KernelBase();
+
+  /// Creates a kernel object from a StreamExecutor and a MultiKernelLoaderSpec.
+  ///
+  /// The StreamExecutor knows which platform it belongs to and the
+  /// MultiKernelLoaderSpec knows how to find the kernel code for different
+  /// platforms, so the combined information is enough to get the kernel code
+  /// for the appropriate platform.
+  static Expected<KernelBase> create(StreamExecutor *ParentExecutor,
+                                     const MultiKernelLoaderSpec &Spec);
+
+  const std::string &getName() const { return Name; }
+  const std::string &getDemangledName() const { return DemangledName; }
+
+  /// Gets a pointer to the platform-specific implementation of this kernel.
+  KernelInterface *getImplementation() { return Implementation.get(); }
+
+private:
+  KernelBase(StreamExecutor *ParentExecutor, const std::string &Name,
+             const std::string &DemangledName,
+             std::unique_ptr<KernelInterface> Implementation);
+
+  StreamExecutor *ParentExecutor;
+  std::string Name;
+  std::string DemangledName;
+  std::unique_ptr<KernelInterface> Implementation;
+
+  KernelBase(const KernelBase &) = delete;
+  KernelBase &operator=(const KernelBase &) = delete;
+};
+
+/// A device kernel function with specified parameter types.
+template <typename... ParameterTs> class TypedKernel : public KernelBase {
+public:
+  TypedKernel(TypedKernel &&) = default;
+  TypedKernel &operator=(TypedKernel &&) = default;
+
+  /// Parameters here have the same meaning as in KernelBase::create.
+  static Expected<TypedKernel> create(StreamExecutor *ParentExecutor,
+                                      const MultiKernelLoaderSpec &Spec) {
+    auto MaybeBase = KernelBase::create(ParentExecutor, Spec);
+    if (!MaybeBase) {
+      return MaybeBase.takeError();
+    }
+    TypedKernel Instance(std::move(*MaybeBase));
+    return std::move(Instance);
+  }
+
+private:
+  TypedKernel(KernelBase &&Base) : KernelBase(std::move(Base)) {}
+
+  TypedKernel(const TypedKernel &) = delete;
+  TypedKernel &operator=(const TypedKernel &) = delete;
+};
+
+} // namespace streamexecutor
+
+#endif // STREAMEXECUTOR_KERNEL_H
diff --git a/parallel-libs/streamexecutor/include/streamexecutor/KernelSpec.h b/parallel-libs/streamexecutor/include/streamexecutor/KernelSpec.h
index a6943143086..110c803258f 100644
--- a/parallel-libs/streamexecutor/include/streamexecutor/KernelSpec.h
+++ b/parallel-libs/streamexecutor/include/streamexecutor/KernelSpec.h
@@ -198,6 +198,13 @@ private:
 /// than doing it by hand.
 class MultiKernelLoaderSpec {
 public:
+  std::string getKernelName() const {
+    if (TheKernelName) {
+      return *TheKernelName;
+    }
+    return "";
+  }
+
   // Convenience getters for testing whether these platform variants have
   // kernel loader specifications available.
 
diff --git a/parallel-libs/streamexecutor/include/streamexecutor/StreamExecutor.h b/parallel-libs/streamexecutor/include/streamexecutor/StreamExecutor.h
new file mode 100644
index 00000000000..b2dd5905797
--- /dev/null
+++ b/parallel-libs/streamexecutor/include/streamexecutor/StreamExecutor.h
@@ -0,0 +1,39 @@
+//===-- StreamExecutor.h - The StreamExecutor class -------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// The StreamExecutor class which represents a single device of a specific
+/// platform.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef STREAMEXECUTOR_STREAMEXECUTOR_H
+#define STREAMEXECUTOR_STREAMEXECUTOR_H
+
+#include "streamexecutor/Utils/Error.h"
+
+namespace streamexecutor {
+
+class KernelInterface;
+
+class StreamExecutor {
+public:
+  /// Gets the kernel implementation for the underlying platform.
+  virtual Expected<std::unique_ptr<KernelInterface>>
+  getKernelImplementation(const MultiKernelLoaderSpec &Spec) {
+    // TODO(jhen): Implement this.
+    return nullptr;
+  }
+
+  // TODO(jhen): Add other methods.
+};
+
+} // namespace streamexecutor
+
+#endif // STREAMEXECUTOR_STREAMEXECUTOR_H
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