[SE] Host platform implementation

[jhen]

Summary

This implementation does not currently support multiple concurrent streams, and
it won't allow kernels to be launched with grids larger than one block or
blocks larger than one thread. These limitations could be removed in the future
by launching new threads on the host, but that is not done in this
implementation.

[jlebar]

streamexecutor/include/streamexecutor/KernelSpec.h

204 ↗

(On Diff #71053)

Is the second const necessary? I feel like it kind of confuses things without really affecting the meaning of the signature in any way. And the second const has no bearing on whether or not you can bind an std::function to a given function: https://godbolt.org/g/XiksZ9

223 ↗

(On Diff #71053)

Nit, maybe HostFunction != nullptr would be clearer? They mean the same thing.

242 ↗

(On Diff #71053)

Why do we return this as a pointer rather than a reference?

streamexecutor/include/streamexecutor/platforms/host/HostPlatform.h

37 ↗

(On Diff #71053)

Isn't the Platform supposed to own the Device object?

streamexecutor/include/streamexecutor/platforms/host/HostPlatformDevice.h

38 ↗

(On Diff #71053)

Oh, I see why you return a pointer from getHostFunction().

Um. You could still return a reference and take the address here... That doesn't actually seem so bad to me.

In either case you probably want to make these multi-kernel objects not-movable, or otherwise allocate the std::function inside a unique_ptr. Because if this thing gets moved we're currently screwed.

151 ↗

(On Diff #71053)

Does this file supercede the for-testing platform we have? (Not necessary to do the switch in this patch, of course.)

[jhen]

• Respond to jlebar's comments

streamexecutor/include/streamexecutor/KernelSpec.h

204 ↗

(On Diff #71053)

You were right, it wasn't necessary. I removed it.

223 ↗

(On Diff #71053)

It's a unique_ptr now and it uses the != nullptr test.

242 ↗

(On Diff #71053)

Changed it to return a reference.

streamexecutor/include/streamexecutor/platforms/host/HostPlatform.h

37 ↗

(On Diff #71053)

Right. I totally forgot about that. Now it does.

streamexecutor/include/streamexecutor/platforms/host/HostPlatformDevice.h

38 ↗

(On Diff #71053)

Now returning a ref and taking the address.

Also converted the MultiKernelLoaderSpec to store a unique_ptr to the function in order to correctly handle moving.

151 ↗

(On Diff #71053)

Yes, I think I will get rid of the for-testing host platform in another patch.

[jlebar]

streamexecutor/include/streamexecutor/KernelSpec.h

273 ↗

(On Diff #71184)

Nit, std::move(Function). (std::function objects may be heavyweight in general, even though in this case it probably doesn't matter.)

streamexecutor/include/streamexecutor/platforms/host/HostPlatform.h

39 ↗

(On Diff #71184)

Do we need this to be thread-safe?

39 ↗

(On Diff #71184)

I'm confused where we call delete on the HostPlatformDevice. It's not super important, but I think if we don't it may show up as a leak in asan, and that means it will show up as a leak in user programs.

[jhen]

• Make thread safe and destroy PlatformDevice

streamexecutor/include/streamexecutor/KernelSpec.h

273 ↗

(On Diff #71184)

I made that change and also changed the signature of the calling function, addHostFunction to take the function argument as an rvalue ref to prevent the copy during the function call.

streamexecutor/include/streamexecutor/platforms/host/HostPlatform.h

39 ↗

(On Diff #71184)

Yes, we do want thread safety. Thanks for catching that. I added a mutex to protect it. I didn't do lazy static initialization because, as jprice suggested before, I think we are going to want to add another method to clear out the platform.

I also made the HostDevicePlatform explicitly owned by HostPlatform so it's clear when it is destroyed.

[jlebar]

streamexecutor/include/streamexecutor/KernelSpec.h

273 ↗

(On Diff #71193)

The usual convention I've seen is to take the heavyweight, copyable arg (be it an std::function, std::string, std::vector, whatever) that you're taking ownership of by value, rather than by rvalue ref.

You're right that making the function arg an rvalue ref prevents accidental copies. The problem is that it also makes *intentional* copies a pain. I think you have to spell out the whole type again:

std::function<const void*> fn = ...;
addHostFunction("kernel_name", std::function<const void*>(fn));

[jhen]

• Change rvalue ref arg to value

parallel-libs/trunk/streamexecutor/examples/CMakeLists.txt

 add_executable(cuda_saxpy_example CUDASaxpy.cpp)
 target_link_libraries(cuda_saxpy_example streamexecutor)
+
+add_executable(host_saxpy_example HostSaxpy.cpp)
+target_link_libraries(host_saxpy_example streamexecutor)

parallel-libs/trunk/streamexecutor/examples/CUDASaxpy.cpp

 /// StreamExecutor API. Snippets of this file will be included as code examples
 /// in documentation. Taking these examples from a real source file guarantees
 /// that the examples will always compile.
 ///
 //===----------------------------------------------------------------------===//
 
 #include <algorithm>
 #include <cassert>
-#include <cstdio>
 #include <cstdlib>
 #include <vector>
 
 #include "streamexecutor/StreamExecutor.h"
 
 /// [Example saxpy compiler-generated]
 // Code in this namespace is generated by the compiler (e.g. clang).
 //

parallel-libs/trunk/streamexecutor/examples/HostSaxpy.cpp

+//===-- HostSaxpy.cpp - Example of host saxpy with StreamExecutor API -----===//
+//
+//                     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 example code demonstrating the usage of the
+/// StreamExecutor API for a host platform.
+///
+//===----------------------------------------------------------------------===//
+
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <vector>
+
+#include "streamexecutor/StreamExecutor.h"
+
+void Saxpy(float A, float *X, float *Y, size_t N) {
+  for (size_t I = 0; I < N; ++I)
+    X[I] = A * X[I] + Y[I];
+}
+
+namespace __compilergen {
+using SaxpyKernel =
+    streamexecutor::Kernel<float, streamexecutor::GlobalDeviceMemory<float>,
+                           streamexecutor::GlobalDeviceMemory<float>, size_t>;
+
+// Wrapper function converts argument addresses to arguments.
+void SaxpyWrapper(const void *const *ArgumentAddresses) {
+  Saxpy(*static_cast<const float *>(ArgumentAddresses[0]),
+        static_cast<float *>(const_cast<void *>(ArgumentAddresses[1])),
+        static_cast<float *>(const_cast<void *>(ArgumentAddresses[2])),
+        *static_cast<const size_t *>(ArgumentAddresses[3]));
+}
+
+// The wrapper function is what gets registered.
+static streamexecutor::MultiKernelLoaderSpec SaxpyLoaderSpec = []() {
+  streamexecutor::MultiKernelLoaderSpec Spec;
+  Spec.addHostFunction("Saxpy", SaxpyWrapper);
+  return Spec;
+}();
+} // namespace __compilergen
+
+int main() {
+  namespace se = ::streamexecutor;
+  namespace cg = ::__compilergen;
+
+  // Create some host data.
+  float A = 42.0f;
+  std::vector<float> HostX = {0, 1, 2, 3};
+  std::vector<float> HostY = {4, 5, 6, 7};
+  size_t ArraySize = HostX.size();
+
+  // Get a device object.
+  se::Platform *Platform =
+      getOrDie(se::PlatformManager::getPlatformByName("host"));
+  if (Platform->getDeviceCount() == 0) {
+    return EXIT_FAILURE;
+  }
+  se::Device *Device = getOrDie(Platform->getDevice(0));
+
+  // Load the kernel onto the device.
+  cg::SaxpyKernel Kernel =
+      getOrDie(Device->createKernel<cg::SaxpyKernel>(cg::SaxpyLoaderSpec));
+
+  se::RegisteredHostMemory<float> RegisteredX =
+      getOrDie(Device->registerHostMemory<float>(HostX));
+  se::RegisteredHostMemory<float> RegisteredY =
+      getOrDie(Device->registerHostMemory<float>(HostY));
+
+  // Allocate memory on the device.
+  se::GlobalDeviceMemory<float> X =
+      getOrDie(Device->allocateDeviceMemory<float>(ArraySize));
+  se::GlobalDeviceMemory<float> Y =
+      getOrDie(Device->allocateDeviceMemory<float>(ArraySize));
+
+  // Run operations on a stream.
+  se::Stream Stream = getOrDie(Device->createStream());
+  Stream.thenCopyH2D(RegisteredX, X)
+      .thenCopyH2D(RegisteredY, Y)
+      .thenLaunch(1, 1, Kernel, A, X, Y, ArraySize)
+      .thenCopyD2H(X, RegisteredX);
+  // Wait for the stream to complete.
+  se::dieIfError(Stream.blockHostUntilDone());
+
+  // Process output data in HostX.
+  std::vector<float> ExpectedX = {4, 47, 90, 133};
+  assert(std::equal(ExpectedX.begin(), ExpectedX.end(), HostX.begin()));
+}

parallel-libs/trunk/streamexecutor/include/streamexecutor/KernelSpec.h

 /// host code to launch the kernel on the device as desired.
 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef STREAMEXECUTOR_KERNELSPEC_H
 #define STREAMEXECUTOR_KERNELSPEC_H
 
 #include <cassert>
+#include <functional>
 #include <map>
 #include <memory>
 #include <string>
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringRef.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

 /// so passing in specs with different kernel names can cause the program to
 /// abort.
 ///
 /// This interface is prone to errors, so it is better to leave
 /// MultiKernelLoaderSpec creation and initialization to the compiler rather
 /// than doing it by hand.
 class MultiKernelLoaderSpec {
 public:
+  /// Type of functions used as host platform kernels.
+  using HostFunctionTy = std::function<void(const void **)>;
+
   std::string getKernelName() const {
     if (TheKernelName)
       return *TheKernelName;
     return "";
   }
 
   // Convenience getters for testing whether these platform variants have
   // kernel loader specifications available.
 
   bool hasCUDAPTXInMemory() const { return TheCUDAPTXInMemorySpec != nullptr; }
   bool hasCUDAFatbinInMemory() const {
     return TheCUDAFatbinInMemorySpec != nullptr;
   }
   bool hasOpenCLTextInMemory() const {
     return TheOpenCLTextInMemorySpec != nullptr;
   }
+  bool hasHostFunction() const { return HostFunction != nullptr; }
 
   // Accessors for platform variant kernel load specifications.
   //
   // Precondition: corresponding has* method returns true.
 
   const CUDAPTXInMemorySpec &getCUDAPTXInMemory() const {
     assert(hasCUDAPTXInMemory() && "getting spec that is not present");
     return *TheCUDAPTXInMemorySpec;
   }
   const CUDAFatbinInMemorySpec &getCUDAFatbinInMemory() const {
     assert(hasCUDAFatbinInMemory() && "getting spec that is not present");
     return *TheCUDAFatbinInMemorySpec;
   }
   const OpenCLTextInMemorySpec &getOpenCLTextInMemory() const {
     assert(hasOpenCLTextInMemory() && "getting spec that is not present");
     return *TheOpenCLTextInMemorySpec;
   }
 
+  const HostFunctionTy &getHostFunction() const {
+    assert(hasHostFunction() && "getting spec that is not present");
+    return *HostFunction;
+  }
+
   // 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 extern "C".
 
   /// Does not take ownership of the PTXCode pointers in the SpecList elements.
   MultiKernelLoaderSpec &
   addCUDAPTXInMemory(llvm::StringRef KernelName,
                      llvm::ArrayRef<CUDAPTXInMemorySpec::PTXSpec> SpecList);
 
   /// Does not take ownership of the FatbinBytes pointer.
   MultiKernelLoaderSpec &addCUDAFatbinInMemory(llvm::StringRef KernelName,
                                                const void *FatbinBytes);
 
   /// Does not take ownership of the OpenCLText pointer.
   MultiKernelLoaderSpec &addOpenCLTextInMemory(llvm::StringRef KernelName,
                                                const char *OpenCLText);
 
+  MultiKernelLoaderSpec &addHostFunction(llvm::StringRef KernelName,
+                                         HostFunctionTy Function) {
+    HostFunction = llvm::make_unique<HostFunctionTy>(std::move(Function));
+    return *this;
+  }
+
 private:
   void setKernelName(llvm::StringRef KernelName);
 
   std::unique_ptr<std::string> TheKernelName;
   std::unique_ptr<CUDAPTXInMemorySpec> TheCUDAPTXInMemorySpec;
   std::unique_ptr<CUDAFatbinInMemorySpec> TheCUDAFatbinInMemorySpec;
   std::unique_ptr<OpenCLTextInMemorySpec> TheOpenCLTextInMemorySpec;
+  std::unique_ptr<HostFunctionTy> HostFunction;
 };
 
 } // namespace streamexecutor
 
 #endif // STREAMEXECUTOR_KERNELSPEC_H

parallel-libs/trunk/streamexecutor/include/streamexecutor/PlatformDevice.h

virtual Error synchronousCopyH2D(const void *HostSrc, size_t SrcByteOffset,
                                    size_t DstByteOffset, size_t ByteCount) {
     return make_error("synchronousCopyH2D not implemented for platform " +
                       getName());
   }
 
   /// Similar to synchronousCopyD2H(const void *, size_t, void
   /// *, size_t, size_t), but copies memory from one location in device memory
   /// to another rather than from device to host.
-  virtual Error synchronousCopyD2D(const void *DeviceDstHandle,
-                                   size_t DstByteOffset,
-                                   const void *DeviceSrcHandle,
-                                   size_t SrcByteOffset, size_t ByteCount) {
+  virtual Error synchronousCopyD2D(const void *DeviceSrcHandle,
+                                   size_t SrcByteOffset,
+                                   const void *DeviceDstHandle,
+                                   size_t DstByteOffset, size_t ByteCount) {
     return make_error("synchronousCopyD2D not implemented for platform " +
                       getName());
   }
 };
 
 } // namespace streamexecutor
 
 #endif // STREAMEXECUTOR_PLATFORMDEVICE_H

parallel-libs/trunk/streamexecutor/include/streamexecutor/platforms/host/HostPlatform.h

+//===-- HostPlatform.h - Host platform subclass -----------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// Declaration of the HostPlatform class.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef STREAMEXECUTOR_PLATFORMS_HOST_HOSTPLATFORM_H
+#define STREAMEXECUTOR_PLATFORMS_HOST_HOSTPLATFORM_H
+
+#include "HostPlatformDevice.h"
+#include "streamexecutor/Device.h"
+#include "streamexecutor/Platform.h"
+
+#include "llvm/Support/Mutex.h"
+
+namespace streamexecutor {
+namespace host {
+
+/// Platform that performs work on the host rather than offloading to an
+/// accelerator.
+class HostPlatform : public Platform {
+public:
+  size_t getDeviceCount() const override { return 1; }
+
+  Expected<Device *> getDevice(size_t DeviceIndex) override {
+    if (DeviceIndex != 0) {
+      return make_error(
+          "Requested device index " + llvm::Twine(DeviceIndex) +
+          " from host platform which only supports device index 0");
+    }
+    llvm::sys::ScopedLock Lock(Mutex);
+    if (!TheDevice) {
+      ThePlatformDevice = llvm::make_unique<HostPlatformDevice>();
+      TheDevice = llvm::make_unique<Device>(ThePlatformDevice.get());
+    }
+    return TheDevice.get();
+  }
+
+private:
+  llvm::sys::Mutex Mutex;
+  std::unique_ptr<HostPlatformDevice> ThePlatformDevice;
+  std::unique_ptr<Device> TheDevice;
+};
+
+} // namespace host
+} // namespace streamexecutor
+
+#endif // STREAMEXECUTOR_PLATFORMS_HOST_HOSTPLATFORM_H

parallel-libs/trunk/streamexecutor/include/streamexecutor/platforms/host/HostPlatformDevice.h

+//===-- HostPlatformDevice.h - HostPlatformDevice class ---------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// Declaration of the HostPlatformDevice class.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef STREAMEXECUTOR_PLATFORMS_HOST_HOSTPLATFORMDEVICE_H
+#define STREAMEXECUTOR_PLATFORMS_HOST_HOSTPLATFORMDEVICE_H
+
+#include <cstdlib>
+#include <cstring>
+
+#include "streamexecutor/PlatformDevice.h"
+
+namespace streamexecutor {
+namespace host {
+
+/// A concrete PlatformDevice subclass that performs its work on the host rather
+/// than offloading to an accelerator.
+class HostPlatformDevice : public PlatformDevice {
+public:
+  std::string getName() const override { return "host"; }
+
+  Expected<const void *>
+  createKernel(const MultiKernelLoaderSpec &Spec) override {
+    if (!Spec.hasHostFunction()) {
+      return make_error("no host implementation available for kernel " +
+                        Spec.getKernelName());
+    }
+    return static_cast<const void *>(&Spec.getHostFunction());
+  }
+
+  Error destroyKernel(const void *Handle) override { return Error::success(); }
+
+  Expected<const void *> createStream() override {
+    // TODO(jhen): Do something with threads to allow multiple streams.
+    return this;
+  }
+
+  Error destroyStream(const void *Handle) override { return Error::success(); }
+
+  Error launch(const void *PlatformStreamHandle, BlockDimensions BlockSize,
+               GridDimensions GridSize, const void *PKernelHandle,
+               const PackedKernelArgumentArrayBase &ArgumentArray) override {
+    // TODO(jhen): Can we do something with BlockSize and GridSize?
+    if (!(BlockSize.X == 1 && BlockSize.Y == 1 && BlockSize.Z == 1)) {
+      return make_error(
+          "Block dimensions were (" + llvm::Twine(BlockSize.X) + "," +
+          llvm::Twine(BlockSize.Y) + "," + llvm::Twine(BlockSize.Z) +
+          "), but only size (1,1,1) is permitted for this platform");
+    }
+    if (!(GridSize.X == 1 && GridSize.Y == 1 && GridSize.Z == 1)) {
+      return make_error(
+          "Grid dimensions were (" + llvm::Twine(GridSize.X) + "," +
+          llvm::Twine(GridSize.Y) + "," + llvm::Twine(GridSize.Z) +
+          "), but only size (1,1,1) is permitted for this platform");
+    }
+
+    (*static_cast<const std::function<void(const void *const *)> *>(
+        PKernelHandle))(ArgumentArray.getAddresses());
+    return Error::success();
+  }
+
+  Error copyD2H(const void *PlatformStreamHandle, const void *DeviceSrcHandle,
+                size_t SrcByteOffset, void *HostDst, size_t DstByteOffset,
+                size_t ByteCount) override {
+    std::memcpy(offset(HostDst, DstByteOffset),
+                offset(DeviceSrcHandle, SrcByteOffset), ByteCount);
+    return Error::success();
+  }
+
+  Error copyH2D(const void *PlatformStreamHandle, const void *HostSrc,
+                size_t SrcByteOffset, const void *DeviceDstHandle,
+                size_t DstByteOffset, size_t ByteCount) override {
+    std::memcpy(offset(DeviceDstHandle, DstByteOffset),
+                offset(HostSrc, SrcByteOffset), ByteCount);
+    return Error::success();
+  }
+
+  Error copyD2D(const void *PlatformStreamHandle, const void *DeviceSrcHandle,
+                size_t SrcByteOffset, const void *DeviceDstHandle,
+                size_t DstByteOffset, size_t ByteCount) override {
+    std::memcpy(offset(DeviceDstHandle, DstByteOffset),
+                offset(DeviceSrcHandle, SrcByteOffset), ByteCount);
+    return Error::success();
+  }
+
+  Error blockHostUntilDone(const void *PlatformStreamHandle) override {
+    // All host operations are synchronous anyway.
+    return Error::success();
+  }
+
+  Expected<void *> allocateDeviceMemory(size_t ByteCount) override {
+    return std::malloc(ByteCount);
+  }
+
+  Error freeDeviceMemory(const void *Handle) override {
+    std::free(const_cast<void *>(Handle));
+    return Error::success();
+  }
+
+  Error registerHostMemory(void *Memory, size_t ByteCount) override {
+    return Error::success();
+  }
+
+  Error unregisterHostMemory(const void *Memory) override {
+    return Error::success();
+  }
+
+  Error synchronousCopyD2H(const void *DeviceSrcHandle, size_t SrcByteOffset,
+                           void *HostDst, size_t DstByteOffset,
+                           size_t ByteCount) override {
+    std::memcpy(offset(HostDst, DstByteOffset),
+                offset(DeviceSrcHandle, SrcByteOffset), ByteCount);
+    return Error::success();
+  }
+
+  Error synchronousCopyH2D(const void *HostSrc, size_t SrcByteOffset,
+                           const void *DeviceDstHandle, size_t DstByteOffset,
+                           size_t ByteCount) override {
+    std::memcpy(offset(DeviceDstHandle, DstByteOffset),
+                offset(HostSrc, SrcByteOffset), ByteCount);
+    return Error::success();
+  }
+
+  Error synchronousCopyD2D(const void *DeviceSrcHandle, size_t SrcByteOffset,
+                           const void *DeviceDstHandle, size_t DstByteOffset,
+                           size_t ByteCount) override {
+    std::memcpy(offset(DeviceDstHandle, DstByteOffset),
+                offset(DeviceSrcHandle, SrcByteOffset), ByteCount);
+    return Error::success();
+  }
+
+private:
+  static void *offset(const void *Base, size_t Offset) {
+    return const_cast<char *>(static_cast<const char *>(Base) + Offset);
+  }
+};
+
+} // namespace host
+} // namespace streamexecutor
+
+#endif // STREAMEXECUTOR_PLATFORMS_HOST_HOSTPLATFORMDEVICE_H

parallel-libs/trunk/streamexecutor/lib/PlatformManager.cpp

 //===-- PlatformManager.cpp - PlatformManager implementation --------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 ///
 /// \file
 /// Implementation of PlatformManager class internals.
 ///
 //===----------------------------------------------------------------------===//
 
 #include "streamexecutor/PlatformManager.h"
+#include "streamexecutor/platforms/host/HostPlatform.h"
 
 namespace streamexecutor {
 
 PlatformManager::PlatformManager() {
   // TODO(jhen): Register known platforms by name.
   // We have a couple of options here:
   //  * Use build-system flags to set preprocessor macros that select the
   //    appropriate code to include here.
   //  * Use static initialization tricks to have platform libraries register
   //    themselves when they are loaded.
+
+  PlatformsByName.emplace("host", llvm::make_unique<host::HostPlatform>());
 }
 
 Expected<Platform *> PlatformManager::getPlatformByName(llvm::StringRef Name) {
   static PlatformManager Instance;
   auto Iterator = Instance.PlatformsByName.find(Name.lower());
   if (Iterator != Instance.PlatformsByName.end())
     return Iterator->second.get();
   return make_error("no available platform with name " + Name);
 }
 
 } // namespace streamexecutor