[rocm-libraries] ROCm/rocm-libraries#6758 (commit 85d3085)

[hipCUB] Update caching device allocator unit test to be more
 reliable (#6758)

## Motivation

The current hipCUB test_caching_device_allocator test was taken from
upstream CUB's corresponding test_allocator.cu. Since then the upstream
CUB test has been updated to be more reliable
(NVIDIA/cccl#5899) but this improvement was not
ported over to hipCUB.

It is possible that the hipCUB test is flaking out at
ROCm/rocm-libraries#6689. This PR hopes to
address any inherent flakiness with the test.

It's possible this change also allows the test to be run under Valgrind
and ASAN; however that is an investigation for another day. I don't want
to delay this PR as it blocking TheRock.

## Technical Details

Borrowing from upstream CUB's solution - implement a mechanism which
launches a kernel which blocks on the device until released. This allows
the test to run deterministically and not rely on a certain luck factor
to succeed.

## Test Plan

Run the test 1000 times on gfx942 in Linux.

## Test Result

All the tests passed.

## Submission Checklist

- [ ] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

Author: stanleytsang-amd

test/hipcub/test_hipcub_caching_device_allocator.cpp

@@ -51,6 +51,55 @@
     #define IS_ASAN_BUILD 0
 #endif
 
+// Borrowing upstream blocking_kernel for host-side control of kernel lifetimes.
+// This kernel does very bad things that violate the HIP programming model, but
+// is okay for the tests here. Do not use this pattern in production code.
+
+// Once launched, this kernel will block the stream until `flag` updates to non-zero.
+__global__ void block_stream(const volatile int32_t* flag)
+{
+	while (!(*flag))
+	{
+	}
+}
+
+struct blocking_kernel
+{
+	blocking_kernel(const hipStream_t& stream)
+	  : m_stream(stream)
+	{
+		HIP_CHECK(hipHostRegister(&m_host_flag, sizeof(m_host_flag), hipHostRegisterMapped));
+		HIP_CHECK(hipHostGetDevicePointer((void**)&m_device_flag, &m_host_flag, 0));
+	}
+	~blocking_kernel()
+	{
+		HIP_CHECK(hipHostUnregister(&m_host_flag));
+	}
+
+	void block()
+	{
+		printf("Blocking Stream %lld\n", (long long) m_stream);
+		m_host_flag = 0;
+		hipLaunchKernelGGL(
+				HIP_KERNEL_NAME(block_stream),
+				dim3(1), dim3(1), 0, m_stream,
+				m_device_flag
+			);		
+	}
+
+	void unblock()
+	{
+		volatile int32_t& flag = m_host_flag;
+		flag                              = 1;
+		printf("Unblocking Stream %lld\n", (long long) m_stream);
+	}
+
+private:
+	int32_t m_host_flag{};
+	int32_t* m_device_flag{};
+	hipStream_t m_stream{0};
+};
+
 __global__
 void EmptyKernel()
 {}
@@ -97,11 +146,9 @@ TEST(HipcubCachingDeviceAllocatorTests, Test1)
     char *d_999B_stream0_b;
     HIP_CHECK(allocator.DeviceAllocate((void **) &d_999B_stream0_a, 999, 0));
 
-    // Run some big kernel in stream 0
-    hipLaunchKernelGGL(
-        HIP_KERNEL_NAME(EmptyKernel),
-        dim3(32000), dim3(256), 1024 * 8, 0
-    );
+	// Run a kernel on stream 0
+	blocking_kernel block_0_a(0);
+	block_0_a.block();
     HIP_CHECK(hipGetLastError());
 
     // Free d_999B_stream0_a
@@ -116,11 +163,9 @@ TEST(HipcubCachingDeviceAllocatorTests, Test1)
     // Check that that we have no cached block on the initial GPU
     ASSERT_EQ(allocator.cached_blocks.size(), 0u);
 
-    // Run some big kernel in stream 0
-    hipLaunchKernelGGL(
-        HIP_KERNEL_NAME(EmptyKernel),
-        dim3(32000), dim3(256), 1024 * 8, 0
-    );
+    // Launch another kernel on stream 0
+	blocking_kernel block_0_b(0);
+	block_0_b.block();	
     HIP_CHECK(hipGetLastError());
 
     // Free d_999B_stream0_b
@@ -137,17 +182,18 @@ TEST(HipcubCachingDeviceAllocatorTests, Test1)
     // Check that that we have one cached block on the initial GPU
     ASSERT_EQ(allocator.cached_blocks.size(), 1u);
 
-    // Run some big kernel in other_stream
-    hipLaunchKernelGGL(
-        HIP_KERNEL_NAME(EmptyKernel),
-        dim3(32000), dim3(256), 1024 * 8, other_stream
-    );
+    // Now run a kernel in other_stream
+	blocking_kernel block_other(other_stream);
+	block_other.block();	
     HIP_CHECK(hipGetLastError());
 
     // Free d_999B_stream_other
     HIP_CHECK(allocator.DeviceFree(d_999B_stream_other_a));
 
-    // Check that we can now use both allocations in stream 0 after synchronizing the device
+    // Check that we can now use both allocations in stream 0 after unblocking both kernels:
+	block_0_a.unblock();
+	block_0_b.unblock();
+	block_other.unblock();	
     HIP_CHECK(hipDeviceSynchronize());
     HIP_CHECK(allocator.DeviceAllocate((void **) &d_999B_stream0_a, 999, 0));
     HIP_CHECK(allocator.DeviceAllocate((void **) &d_999B_stream0_b, 999, 0));
@@ -174,18 +220,16 @@ TEST(HipcubCachingDeviceAllocatorTests, Test1)
     ASSERT_EQ(allocator.cached_blocks.size(), 0u);
 
     // Run some big kernel in other_stream
-    hipLaunchKernelGGL(
-        HIP_KERNEL_NAME(EmptyKernel),
-        dim3(32000), dim3(256), 1024 * 8, other_stream
-    );
+	block_other.block();	
     HIP_CHECK(hipGetLastError());
 
     // Free d_999B_stream_other_a and d_999B_stream_other_b
     HIP_CHECK(allocator.DeviceFree(d_999B_stream_other_a));
     HIP_CHECK(allocator.DeviceFree(d_999B_stream_other_b));
 
     // Check that we can now use both allocations in stream 0 after synchronizing the device and destroying the other stream
-    HIP_CHECK(hipDeviceSynchronize());
+    block_other.unblock();
+	HIP_CHECK(hipDeviceSynchronize());
     HIP_CHECK(hipStreamDestroy(other_stream));
     HIP_CHECK(allocator.DeviceAllocate((void **) &d_999B_stream0_a, 999, 0));
     HIP_CHECK(allocator.DeviceAllocate((void **) &d_999B_stream0_b, 999, 0));