[SE] Pack global dev handle addresses

[jhen]

Summary

We were packing global device memory handles in
PackedKernelArgumentArray, but as I was implementing the CUDA
platform, I realized that CUDA wants the address of the handle, not the
handle itself. So this patch switches to packing the address of the
handle.

[jlebar]

streamexecutor/include/streamexecutor/DeviceMemory.h

170

Same question as last patch: What happens if this guy is moved?

Specifically, would it be a problem if it were moved after calling thenLaunch but before the driver actually launches the kernel?

[jhen]

• Warning about kernel launch args

streamexecutor/include/streamexecutor/DeviceMemory.h

170

Yes, it would be a problem if the argument was moved at the wrong time. This was a choice we made internally to reduce the kernel launch overhead. Apparently it could make up 5% of some applications' run-time.

In response to your question, I added a warning message to the Stream::thenLaunch method. This message tells users not to touch kernel launch arguments from other threads.

[jlebar]

streamexecutor/include/streamexecutor/DeviceMemory.h

170

Clearly in the multithreaded case you can't modify a GlobalDeviceMemoryBase concurrently with a call to thenLaunch(). That's true of any parameters to any function call, so I am not sure that's even worth warning about.

What I was worried about was a single-threaded case, something like this:

GlobalDeviceMemory<int> other_mem;
{  
  GlobalDeviceMemory<int> mem;
  Stream->thenLaunch(foo_kernel, mem);
  other_mem = std::move(mem);
}
Stream->block();

Is this safe? That is, do we use &mem.Handle only within thenLaunch? (We don't actually have to have a separate scope and so on for us to hit this same problem.)

(I have to admit, if any of this is 5% of some applications' runtime, it seems like we could do a lot better even than what we have here. I'm not sure how, but hot code is hot...)

[jhen]

• Remove warning

streamexecutor/include/streamexecutor/DeviceMemory.h

170

Oh, that's right. I think we spoke about this before. There should be no problem with your example, the platform should be responsible for copying the arguments (using memcpy) before returning from a launch call. This will match the guarantee currently provided by the CUDA driver library.

I removed the unnecessary warning.

[jlebar]

I think we spoke about this before.

Oh, I remember this now. :) Okay, sgtm.

streamexecutor/examples/HostSaxpy.cpp

 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<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;

streamexecutor/include/streamexecutor/DeviceMemory.h

 ///
 /// For example, in the OpenCL platform, the handle is a pointer to a _cl_mem
 /// handle object which really is completely opaque to the user.
 class GlobalDeviceMemoryBase {
 public:
   /// Returns an opaque handle to the underlying memory.
   const void *getHandle() const { return Handle; }
 
+  /// Returns the address of the opaque handle as stored by this object.
+  const void *const *getHandleAddress() const { return &Handle; }
+
   // Cannot copy because the handle must be owned by a single object.
   GlobalDeviceMemoryBase(const GlobalDeviceMemoryBase &) = delete;
   GlobalDeviceMemoryBase &operator=(const GlobalDeviceMemoryBase &) = delete;
 
 protected:
   /// Creates a GlobalDeviceMemoryBase from a handle and a byte count.
   GlobalDeviceMemoryBase(Device *D, const void *Handle, size_t ByteCount)
       : TheDevice(D), Handle(Handle), ByteCount(ByteCount) {}
GlobalDeviceMemoryBase &operator=(GlobalDeviceMemoryBase &&Other) noexcept {
     Other.Handle = nullptr;
     Other.ByteCount = 0;
     return *this;
   }
 
   ~GlobalDeviceMemoryBase();
 
   Device *TheDevice;  // Pointer to the device on which this memory lives.
   const void *Handle; // Platform-dependent value representing allocated memory.
   size_t ByteCount;   // Size in bytes of this allocation.
 };
 
 /// Typed wrapper around the "void *"-like GlobalDeviceMemoryBase class.
 ///
 /// For example, GlobalDeviceMemory<int> is a simple wrapper around
 /// GlobalDeviceMemoryBase that represents a buffer of integers stored in global
 /// device memory.

streamexecutor/include/streamexecutor/PackedKernelArgumentArray.h

private:
 
   // Pack a normal, non-device-memory argument.
   template <typename T> void PackOneArgument(size_t Index, const T &Argument) {
     Addresses[Index] = &Argument;
     Sizes[Index] = sizeof(T);
     Types[Index] = KernelArgumentType::VALUE;
   }
 
-  // Pack a GlobalDeviceMemoryBase argument.
-  void PackOneArgument(size_t Index, const GlobalDeviceMemoryBase &Argument) {
-    Addresses[Index] = Argument.getHandle();
-    Sizes[Index] = sizeof(void *);
-    Types[Index] = KernelArgumentType::GLOBAL_DEVICE_MEMORY;
-  }
-
-  // Pack a GlobalDeviceMemoryBase pointer argument.
-  void PackOneArgument(size_t Index, GlobalDeviceMemoryBase *Argument) {
-    Addresses[Index] = Argument->getHandle();
-    Sizes[Index] = sizeof(void *);
-    Types[Index] = KernelArgumentType::GLOBAL_DEVICE_MEMORY;
-  }
-
-  // Pack a const GlobalDeviceMemoryBase pointer argument.
-  void PackOneArgument(size_t Index, const GlobalDeviceMemoryBase *Argument) {
-    Addresses[Index] = Argument->getHandle();
-    Sizes[Index] = sizeof(void *);
-    Types[Index] = KernelArgumentType::GLOBAL_DEVICE_MEMORY;
-  }
-
   // Pack a GlobalDeviceMemory<T> argument.
   template <typename T>
   void PackOneArgument(size_t Index, const GlobalDeviceMemory<T> &Argument) {
-    Addresses[Index] = Argument.getHandle();
+    Addresses[Index] = Argument.getHandleAddress();
     Sizes[Index] = sizeof(void *);
     Types[Index] = KernelArgumentType::GLOBAL_DEVICE_MEMORY;
   }
 
   // Pack a GlobalDeviceMemory<T> pointer argument.
   template <typename T>
   void PackOneArgument(size_t Index, GlobalDeviceMemory<T> *Argument) {
-    Addresses[Index] = Argument->getHandle();
+    Addresses[Index] = Argument->getHandleAddress();
     Sizes[Index] = sizeof(void *);
     Types[Index] = KernelArgumentType::GLOBAL_DEVICE_MEMORY;
   }
 
   // Pack a const GlobalDeviceMemory<T> pointer argument.
   template <typename T>
   void PackOneArgument(size_t Index, const GlobalDeviceMemory<T> *Argument) {
-    Addresses[Index] = Argument->getHandle();
+    Addresses[Index] = Argument->getHandleAddress();
     Sizes[Index] = sizeof(void *);
     Types[Index] = KernelArgumentType::GLOBAL_DEVICE_MEMORY;
   }
 
   // Pack a SharedDeviceMemory argument.
   template <typename T>
   void PackOneArgument(size_t Index, const SharedDeviceMemory<T> &Argument) {
     ++SharedCount;

streamexecutor/unittests/CoreTests/PackedKernelArgumentArrayTest.cpp

TEST_F(DeviceMemoryPackingTest, SingleValue) {
   auto Array = se::make_kernel_argument_pack(Value);
   ExpectEqual(&Value, sizeof(Value), Type::VALUE, Array, 0);
   EXPECT_EQ(1u, Array.getArgumentCount());
   EXPECT_EQ(0u, Array.getSharedCount());
 }
 
 TEST_F(DeviceMemoryPackingTest, SingleTypedGlobal) {
   auto Array = se::make_kernel_argument_pack(TypedGlobal);
-  ExpectEqual(TypedGlobal.getHandle(), sizeof(void *),
+  ExpectEqual(TypedGlobal.getHandleAddress(), sizeof(void *),
               Type::GLOBAL_DEVICE_MEMORY, Array, 0);
   EXPECT_EQ(1u, Array.getArgumentCount());
   EXPECT_EQ(0u, Array.getSharedCount());
 }
 
 TEST_F(DeviceMemoryPackingTest, SingleTypedGlobalPointer) {
   auto Array = se::make_kernel_argument_pack(&TypedGlobal);
-  ExpectEqual(TypedGlobal.getHandle(), sizeof(void *),
+  ExpectEqual(TypedGlobal.getHandleAddress(), sizeof(void *),
               Type::GLOBAL_DEVICE_MEMORY, Array, 0);
   EXPECT_EQ(1u, Array.getArgumentCount());
   EXPECT_EQ(0u, Array.getSharedCount());
 }
 
 TEST_F(DeviceMemoryPackingTest, SingleConstTypedGlobalPointer) {
   const se::GlobalDeviceMemory<int> *ArgumentPointer = &TypedGlobal;
   auto Array = se::make_kernel_argument_pack(ArgumentPointer);
-  ExpectEqual(TypedGlobal.getHandle(), sizeof(void *),
+  ExpectEqual(TypedGlobal.getHandleAddress(), sizeof(void *),
               Type::GLOBAL_DEVICE_MEMORY, Array, 0);
   EXPECT_EQ(1u, Array.getArgumentCount());
   EXPECT_EQ(0u, Array.getSharedCount());
 }
 
 TEST_F(DeviceMemoryPackingTest, SingleTypedShared) {
   auto Array = se::make_kernel_argument_pack(TypedShared);
   ExpectEqual(nullptr, TypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,

 
 TEST_F(DeviceMemoryPackingTest, PackSeveralArguments) {
   const se::GlobalDeviceMemory<int> *TypedGlobalPointer = &TypedGlobal;
   const se::SharedDeviceMemory<int> *TypedSharedPointer = &TypedShared;
   auto Array = se::make_kernel_argument_pack(Value, TypedGlobal, &TypedGlobal,
                                              TypedGlobalPointer, TypedShared,
                                              &TypedShared, TypedSharedPointer);
   ExpectEqual(&Value, sizeof(Value), Type::VALUE, Array, 0);
-  ExpectEqual(TypedGlobal.getHandle(), sizeof(void *),
+  ExpectEqual(TypedGlobal.getHandleAddress(), sizeof(void *),
               Type::GLOBAL_DEVICE_MEMORY, Array, 1);
-  ExpectEqual(TypedGlobal.getHandle(), sizeof(void *),
+  ExpectEqual(TypedGlobal.getHandleAddress(), sizeof(void *),
               Type::GLOBAL_DEVICE_MEMORY, Array, 2);
-  ExpectEqual(TypedGlobal.getHandle(), sizeof(void *),
+  ExpectEqual(TypedGlobal.getHandleAddress(), sizeof(void *),
               Type::GLOBAL_DEVICE_MEMORY, Array, 3);
   ExpectEqual(nullptr, TypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
               Array, 4);
   ExpectEqual(nullptr, TypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
               Array, 5);
   ExpectEqual(nullptr, TypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
               Array, 6);
   EXPECT_EQ(7u, Array.getArgumentCount());
   EXPECT_EQ(3u, Array.getSharedCount());
 }
 
 } // namespace