streamexecutor/include/streamexecutor/Stream.h - llvm-project/parallel-libs - Git at Google

 //===-- Stream.h - A stream of execution ------------------------*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 ///
 /// \file
 ///
 /// A Stream instance represents a queue of sequential, host-asynchronous work
 /// to be performed on a device.
 ///
 /// To enqueue work on a device, first create a StreamExecutor instance for a
 /// given device and then use that StreamExecutor to create a Stream instance.
 /// The Stream instance will perform its work on the device managed by the
 /// StreamExecutor that created it.
 ///
 /// The various "then" methods of the Stream object, such as thenMemcpyH2D and
 /// thenLaunch, may be used to enqueue work on the Stream, and the
 /// blockHostUntilDone() method may be used to block the host code until the
 /// Stream has completed all its work.
 ///
 /// Multiple Stream instances can be created for the same StreamExecutor. This
 /// allows several independent streams of computation to be performed
 /// simultaneously on a single device.
 ///
 //===----------------------------------------------------------------------===//

 #ifndef STREAMEXECUTOR_STREAM_H
 #define STREAMEXECUTOR_STREAM_H

 #include <cassert>
 #include <memory>
 #include <string>

 #include "streamexecutor/DeviceMemory.h"
 #include "streamexecutor/Kernel.h"
 #include "streamexecutor/LaunchDimensions.h"
 #include "streamexecutor/PackedKernelArgumentArray.h"
 #include "streamexecutor/PlatformInterfaces.h"
 #include "streamexecutor/Utils/Error.h"

 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/RWMutex.h"

 namespace streamexecutor {

 /// Represents a stream of dependent computations on a device.
 ///
 /// The operations within a stream execute sequentially and asynchronously until
 /// blockHostUntilDone() is invoked, which synchronously joins host code with
 /// the execution of the stream.
 ///
 /// If any given operation fails when entraining work for the stream, isOK()
 /// will indicate that an error has occurred and getStatus() will get the first
 /// error that occurred on the stream. There is no way to clear the error state
 /// of a stream once it is in an error state.
 class Stream {
 public:
   explicit Stream(std::unique_ptr<PlatformStreamHandle> PStream);

   ~Stream();

   /// Returns whether any error has occurred while entraining work on this
   /// stream.
   bool isOK() const {
     llvm::sys::ScopedReader ReaderLock(ErrorMessageMutex);
     return !ErrorMessage;
   }

   /// Returns the status created by the first error that occurred while
   /// entraining work on this stream.
   Error getStatus() const {
     llvm::sys::ScopedReader ReaderLock(ErrorMessageMutex);
     if (ErrorMessage)
       return make_error(*ErrorMessage);
     else
       return Error::success();
   };

   /// Entrains onto the stream of operations a kernel launch with the given
   /// arguments.
   ///
   /// These arguments can be device memory types like GlobalDeviceMemory<T> and
   /// SharedDeviceMemory<T>, or they can be primitive types such as int. The
   /// allowable argument types are determined by the template parameters to the
   /// TypedKernel argument.
   template <typename... ParameterTs>
   Stream &thenLaunch(BlockDimensions BlockSize, GridDimensions GridSize,
                      const TypedKernel<ParameterTs...> &Kernel,
                      const ParameterTs &... Arguments) {
     auto ArgumentArray =
         make_kernel_argument_pack<ParameterTs...>(Arguments...);
     setError(PlatformExecutor->launch(ThePlatformStream.get(), BlockSize,
                                       GridSize, Kernel, ArgumentArray));
     return *this;
   }

   /// Entrain onto the stream a memcpy of a given number of elements from a
   /// device source to a host destination.
   ///
   /// HostDst must be a pointer to host memory allocated by
   /// StreamExecutor::allocateHostMemory or otherwise allocated and then
   /// registered with StreamExecutor::registerHostMemory.
   template <typename T>
   Stream &thenMemcpyD2H(const GlobalDeviceMemory<T> &DeviceSrc,
                         llvm::MutableArrayRef<T> HostDst, size_t ElementCount) {
     if (ElementCount > DeviceSrc.getElementCount())
       setError("copying too many elements, " + llvm::Twine(ElementCount) +
                ", from device memory array of size " +
                llvm::Twine(DeviceSrc.getElementCount()));
     else if (ElementCount > HostDst.size())
       setError("copying too many elements, " + llvm::Twine(ElementCount) +
                ", to host array of size " + llvm::Twine(HostDst.size()));
     else
       setError(PlatformExecutor->memcpyD2H(ThePlatformStream.get(), DeviceSrc,
                                            HostDst.data(),
                                            ElementCount * sizeof(T)));
     return *this;
   }

   /// Same as thenMemcpyD2H above, but copies the entire source to the
   /// destination.
   template <typename T>
   Stream &thenMemcpyD2H(const GlobalDeviceMemory<T> &DeviceSrc,
                         llvm::MutableArrayRef<T> HostDst) {
     return thenMemcpyD2H(DeviceSrc, HostDst, DeviceSrc.getElementCount());
   }

   /// Entrain onto the stream a memcpy of a given number of elements from a host
   /// source to a device destination.
   ///
   /// HostSrc must be a pointer to host memory allocated by
   /// StreamExecutor::allocateHostMemory or otherwise allocated and then
   /// registered with StreamExecutor::registerHostMemory.
   template <typename T>
   Stream &thenMemcpyH2D(llvm::ArrayRef<T> HostSrc,
                         GlobalDeviceMemory<T> *DeviceDst, size_t ElementCount) {
     if (ElementCount > HostSrc.size())
       setError("copying too many elements, " + llvm::Twine(ElementCount) +
                ", from host array of size " + llvm::Twine(HostSrc.size()));
     else if (ElementCount > DeviceDst->getElementCount())
       setError("copying too many elements, " + llvm::Twine(ElementCount) +
                ", to device memory array of size " +
                llvm::Twine(DeviceDst->getElementCount()));
     else
       setError(PlatformExecutor->memcpyH2D(ThePlatformStream.get(),
                                            HostSrc.data(), DeviceDst,
                                            ElementCount * sizeof(T)));
     return *this;
   }

   /// Same as thenMemcpyH2D above, but copies the entire source to the
   /// destination.
   template <typename T>
   Stream &thenMemcpyH2D(llvm::ArrayRef<T> HostSrc,
                         GlobalDeviceMemory<T> *DeviceDst) {
     return thenMemcpyH2D(HostSrc, DeviceDst, HostSrc.size());
   }

   /// Entrain onto the stream a memcpy of a given number of elements from a
   /// device source to a device destination.
   template <typename T>
   Stream &thenMemcpyD2D(const GlobalDeviceMemory<T> &DeviceSrc,
                         GlobalDeviceMemory<T> *DeviceDst, size_t ElementCount) {
     if (ElementCount > DeviceSrc.getElementCount())
       setError("copying too many elements, " + llvm::Twine(ElementCount) +
                ", from device memory array of size " +
                llvm::Twine(DeviceSrc.getElementCount()));
     else if (ElementCount > DeviceDst->getElementCount())
       setError("copying too many elements, " + llvm::Twine(ElementCount) +
                ", to device memory array of size " +
                llvm::Twine(DeviceDst->getElementCount()));
     else
       setError(PlatformExecutor->memcpyD2D(ThePlatformStream.get(), DeviceSrc,
                                            DeviceDst,
                                            ElementCount * sizeof(T)));
     return *this;
   }

   /// Same as thenMemcpyD2D above, but copies the entire source to the
   /// destination.
   template <typename T>
   Stream &thenMemcpyD2D(const GlobalDeviceMemory<T> &DeviceSrc,
                         GlobalDeviceMemory<T> *DeviceDst) {
     return thenMemcpyD2D(DeviceSrc, DeviceDst, DeviceSrc.getElementCount());
   }

   /// Blocks the host code, waiting for the operations entrained on the stream
   /// (enqueued up to this point in program execution) to complete.
   ///
   /// Returns true if there are no errors on the stream.
   bool blockHostUntilDone() {
     Error E = PlatformExecutor->blockHostUntilDone(ThePlatformStream.get());
     bool returnValue = static_cast<bool>(E);
     setError(std::move(E));
     return returnValue;
   }

 private:
   /// Sets the error state from an Error object.
   ///
   /// Does not overwrite the error if it is already set.
   void setError(Error &&E) {
     if (E) {
       llvm::sys::ScopedWriter WriterLock(ErrorMessageMutex);
       if (!ErrorMessage)
         ErrorMessage = consumeAndGetMessage(std::move(E));
     }
   }

   /// Sets the error state from an error message.
   ///
   /// Does not overwrite the error if it is already set.
   void setError(llvm::Twine Message) {
     llvm::sys::ScopedWriter WriterLock(ErrorMessageMutex);
     if (!ErrorMessage)
       ErrorMessage = Message.str();
   }

   /// The PlatformStreamExecutor that supports the operations of this stream.
   PlatformStreamExecutor *PlatformExecutor;

   /// The platform-specific stream handle for this instance.
   std::unique_ptr<PlatformStreamHandle> ThePlatformStream;

   /// Mutex that guards the error state flags.
   ///
   /// Mutable so that it can be obtained via const reader lock.
   mutable llvm::sys::RWMutex ErrorMessageMutex;

   /// First error message for an operation in this stream or empty if there have
   /// been no errors.
   llvm::Optional<std::string> ErrorMessage;

   Stream(const Stream &) = delete;
   void operator=(const Stream &) = delete;
 };

 } // namespace streamexecutor

 #endif // STREAMEXECUTOR_STREAM_H
	//===-- Stream.h - A stream of execution ------------------------- C++ --===//
	//
	// The LLVM Compiler Infrastructure
	//
	// This file is distributed under the University of Illinois Open Source
	// License. See LICENSE.TXT for details.
	//
	//===----------------------------------------------------------------------===//
	///
	/// \file
	///
	/// A Stream instance represents a queue of sequential, host-asynchronous work
	/// to be performed on a device.
	///
	/// To enqueue work on a device, first create a StreamExecutor instance for a
	/// given device and then use that StreamExecutor to create a Stream instance.
	/// The Stream instance will perform its work on the device managed by the
	/// StreamExecutor that created it.
	///
	/// The various "then" methods of the Stream object, such as thenMemcpyH2D and
	/// thenLaunch, may be used to enqueue work on the Stream, and the
	/// blockHostUntilDone() method may be used to block the host code until the
	/// Stream has completed all its work.
	///
	/// Multiple Stream instances can be created for the same StreamExecutor. This
	/// allows several independent streams of computation to be performed
	/// simultaneously on a single device.
	///
	//===----------------------------------------------------------------------===//

	#ifndef STREAMEXECUTOR_STREAM_H
	#define STREAMEXECUTOR_STREAM_H

	#include <cassert>
	#include <memory>
	#include <string>

	#include "streamexecutor/DeviceMemory.h"
	#include "streamexecutor/Kernel.h"
	#include "streamexecutor/LaunchDimensions.h"
	#include "streamexecutor/PackedKernelArgumentArray.h"
	#include "streamexecutor/PlatformInterfaces.h"
	#include "streamexecutor/Utils/Error.h"

	#include "llvm/ADT/Optional.h"
	#include "llvm/ADT/Twine.h"
	#include "llvm/Support/RWMutex.h"

	namespace streamexecutor {

	/// Represents a stream of dependent computations on a device.
	///
	/// The operations within a stream execute sequentially and asynchronously until
	/// blockHostUntilDone() is invoked, which synchronously joins host code with
	/// the execution of the stream.
	///
	/// If any given operation fails when entraining work for the stream, isOK()
	/// will indicate that an error has occurred and getStatus() will get the first
	/// error that occurred on the stream. There is no way to clear the error state
	/// of a stream once it is in an error state.
	class Stream {
	public:
	explicit Stream(std::unique_ptr<PlatformStreamHandle> PStream);

	~Stream();

	/// Returns whether any error has occurred while entraining work on this
	/// stream.
	bool isOK() const {
	llvm::sys::ScopedReader ReaderLock(ErrorMessageMutex);
	return !ErrorMessage;
	}

	/// Returns the status created by the first error that occurred while
	/// entraining work on this stream.
	Error getStatus() const {
	llvm::sys::ScopedReader ReaderLock(ErrorMessageMutex);
	if (ErrorMessage)
	return make_error(*ErrorMessage);
	else
	return Error::success();
	};

	/// Entrains onto the stream of operations a kernel launch with the given
	/// arguments.
	///
	/// These arguments can be device memory types like GlobalDeviceMemory<T> and
	/// SharedDeviceMemory<T>, or they can be primitive types such as int. The
	/// allowable argument types are determined by the template parameters to the
	/// TypedKernel argument.
	template <typename... ParameterTs>
	Stream &thenLaunch(BlockDimensions BlockSize, GridDimensions GridSize,
	const TypedKernel<ParameterTs...> &Kernel,
	const ParameterTs &... Arguments) {
	auto ArgumentArray =
	make_kernel_argument_pack<ParameterTs...>(Arguments...);
	setError(PlatformExecutor->launch(ThePlatformStream.get(), BlockSize,
	GridSize, Kernel, ArgumentArray));
	return *this;
	}

	/// Entrain onto the stream a memcpy of a given number of elements from a
	/// device source to a host destination.
	///
	/// HostDst must be a pointer to host memory allocated by
	/// StreamExecutor::allocateHostMemory or otherwise allocated and then
	/// registered with StreamExecutor::registerHostMemory.
	template <typename T>
	Stream &thenMemcpyD2H(const GlobalDeviceMemory<T> &DeviceSrc,
	llvm::MutableArrayRef<T> HostDst, size_t ElementCount) {
	if (ElementCount > DeviceSrc.getElementCount())
	setError("copying too many elements, " + llvm::Twine(ElementCount) +
	", from device memory array of size " +
	llvm::Twine(DeviceSrc.getElementCount()));
	else if (ElementCount > HostDst.size())
	setError("copying too many elements, " + llvm::Twine(ElementCount) +
	", to host array of size " + llvm::Twine(HostDst.size()));
	else
	setError(PlatformExecutor->memcpyD2H(ThePlatformStream.get(), DeviceSrc,
	HostDst.data(),
	ElementCount * sizeof(T)));
	return *this;
	}

	/// Same as thenMemcpyD2H above, but copies the entire source to the
	/// destination.
	template <typename T>
	Stream &thenMemcpyD2H(const GlobalDeviceMemory<T> &DeviceSrc,
	llvm::MutableArrayRef<T> HostDst) {
	return thenMemcpyD2H(DeviceSrc, HostDst, DeviceSrc.getElementCount());
	}

	/// Entrain onto the stream a memcpy of a given number of elements from a host
	/// source to a device destination.
	///
	/// HostSrc must be a pointer to host memory allocated by
	/// StreamExecutor::allocateHostMemory or otherwise allocated and then
	/// registered with StreamExecutor::registerHostMemory.
	template <typename T>
	Stream &thenMemcpyH2D(llvm::ArrayRef<T> HostSrc,
	GlobalDeviceMemory<T> *DeviceDst, size_t ElementCount) {
	if (ElementCount > HostSrc.size())
	setError("copying too many elements, " + llvm::Twine(ElementCount) +
	", from host array of size " + llvm::Twine(HostSrc.size()));
	else if (ElementCount > DeviceDst->getElementCount())
	setError("copying too many elements, " + llvm::Twine(ElementCount) +
	", to device memory array of size " +
	llvm::Twine(DeviceDst->getElementCount()));
	else
	setError(PlatformExecutor->memcpyH2D(ThePlatformStream.get(),
	HostSrc.data(), DeviceDst,
	ElementCount * sizeof(T)));
	return *this;
	}

	/// Same as thenMemcpyH2D above, but copies the entire source to the
	/// destination.
	template <typename T>
	Stream &thenMemcpyH2D(llvm::ArrayRef<T> HostSrc,
	GlobalDeviceMemory<T> *DeviceDst) {
	return thenMemcpyH2D(HostSrc, DeviceDst, HostSrc.size());
	}

	/// Entrain onto the stream a memcpy of a given number of elements from a
	/// device source to a device destination.
	template <typename T>
	Stream &thenMemcpyD2D(const GlobalDeviceMemory<T> &DeviceSrc,
	GlobalDeviceMemory<T> *DeviceDst, size_t ElementCount) {
	if (ElementCount > DeviceSrc.getElementCount())
	setError("copying too many elements, " + llvm::Twine(ElementCount) +
	", from device memory array of size " +
	llvm::Twine(DeviceSrc.getElementCount()));
	else if (ElementCount > DeviceDst->getElementCount())
	setError("copying too many elements, " + llvm::Twine(ElementCount) +
	", to device memory array of size " +
	llvm::Twine(DeviceDst->getElementCount()));
	else
	setError(PlatformExecutor->memcpyD2D(ThePlatformStream.get(), DeviceSrc,
	DeviceDst,
	ElementCount * sizeof(T)));
	return *this;
	}

	/// Same as thenMemcpyD2D above, but copies the entire source to the
	/// destination.
	template <typename T>
	Stream &thenMemcpyD2D(const GlobalDeviceMemory<T> &DeviceSrc,
	GlobalDeviceMemory<T> *DeviceDst) {
	return thenMemcpyD2D(DeviceSrc, DeviceDst, DeviceSrc.getElementCount());
	}

	/// Blocks the host code, waiting for the operations entrained on the stream
	/// (enqueued up to this point in program execution) to complete.
	///
	/// Returns true if there are no errors on the stream.
	bool blockHostUntilDone() {
	Error E = PlatformExecutor->blockHostUntilDone(ThePlatformStream.get());
	bool returnValue = static_cast<bool>(E);
	setError(std::move(E));
	return returnValue;
	}

	private:
	/// Sets the error state from an Error object.
	///
	/// Does not overwrite the error if it is already set.
	void setError(Error &&E) {
	if (E) {
	llvm::sys::ScopedWriter WriterLock(ErrorMessageMutex);
	if (!ErrorMessage)
	ErrorMessage = consumeAndGetMessage(std::move(E));
	}
	}

	/// Sets the error state from an error message.
	///
	/// Does not overwrite the error if it is already set.
	void setError(llvm::Twine Message) {
	llvm::sys::ScopedWriter WriterLock(ErrorMessageMutex);
	if (!ErrorMessage)
	ErrorMessage = Message.str();
	}

	/// The PlatformStreamExecutor that supports the operations of this stream.
	PlatformStreamExecutor *PlatformExecutor;

	/// The platform-specific stream handle for this instance.
	std::unique_ptr<PlatformStreamHandle> ThePlatformStream;

	/// Mutex that guards the error state flags.
	///
	/// Mutable so that it can be obtained via const reader lock.
	mutable llvm::sys::RWMutex ErrorMessageMutex;

	/// First error message for an operation in this stream or empty if there have
	/// been no errors.
	llvm::Optional<std::string> ErrorMessage;

	Stream(const Stream &) = delete;
	void operator=(const Stream &) = delete;
	};

	} // namespace streamexecutor

	#endif // STREAMEXECUTOR_STREAM_H