[SE] Pack global dev handle addresses
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.
Reviewers: jlebar
Subscribers: jprice, jlebar, parallel_libs-commits
Differential Revision: https://reviews.llvm.org/D24528
llvm-svn: 281424
GitOrigin-RevId: b38d8a3a3baabf759e819fdefd764462691f4048
diff --git a/streamexecutor/examples/HostSaxpy.cpp b/streamexecutor/examples/HostSaxpy.cpp
index 5bcbcc8..cf81b0b 100644
--- a/streamexecutor/examples/HostSaxpy.cpp
+++ b/streamexecutor/examples/HostSaxpy.cpp
@@ -33,8 +33,8 @@
// 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]));
}
diff --git a/streamexecutor/include/streamexecutor/DeviceMemory.h b/streamexecutor/include/streamexecutor/DeviceMemory.h
index d8f7cef..62f6e57 100644
--- a/streamexecutor/include/streamexecutor/DeviceMemory.h
+++ b/streamexecutor/include/streamexecutor/DeviceMemory.h
@@ -133,6 +133,9 @@
/// 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;
diff --git a/streamexecutor/include/streamexecutor/PackedKernelArgumentArray.h b/streamexecutor/include/streamexecutor/PackedKernelArgumentArray.h
index ba53ea4..f34ec67 100644
--- a/streamexecutor/include/streamexecutor/PackedKernelArgumentArray.h
+++ b/streamexecutor/include/streamexecutor/PackedKernelArgumentArray.h
@@ -164,31 +164,10 @@
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;
}
@@ -196,7 +175,7 @@
// 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;
}
@@ -204,7 +183,7 @@
// 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;
}
diff --git a/streamexecutor/unittests/CoreTests/PackedKernelArgumentArrayTest.cpp b/streamexecutor/unittests/CoreTests/PackedKernelArgumentArrayTest.cpp
index 443cdde..860f21c 100644
--- a/streamexecutor/unittests/CoreTests/PackedKernelArgumentArrayTest.cpp
+++ b/streamexecutor/unittests/CoreTests/PackedKernelArgumentArrayTest.cpp
@@ -76,7 +76,7 @@
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());
@@ -84,7 +84,7 @@
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());
@@ -93,7 +93,7 @@
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());
@@ -131,11 +131,11 @@
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);
