From 022651b154715a1c2c5372657b4fec5071e479b4 Mon Sep 17 00:00:00 2001
From: Georgy Evtushenko <evtushenko.georgy@gmail.com>
Date: Thu, 8 Jun 2023 15:20:35 +0400
Subject: [PATCH] Workaround three-way partition compilation issue

---
 test/test_device_three_way_partition.cu | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/test/test_device_three_way_partition.cu b/test/test_device_three_way_partition.cu
index dfaab8d300..27fdf00591 100644
--- a/test/test_device_three_way_partition.cu
+++ b/test/test_device_three_way_partition.cu
@@ -379,16 +379,16 @@ template <typename ValueT>
 struct CountToPair
 {
   template <typename OffsetT>
-  __device__ __host__ Pair<ValueT, std::uint64_t>operator()(OffsetT id)
+  __device__ __host__ Pair<ValueT, std::uint32_t> operator()(OffsetT id)
   {
-    return Pair<ValueT, std::uint64_t>(static_cast<ValueT>(id), id);
+    return Pair<ValueT, std::uint32_t>(static_cast<ValueT>(id), id);
   }
 };
 
 template <typename KeyT>
 void TestStability(int num_items)
 {
-  using T = Pair<KeyT, std::uint64_t>;
+  using T = Pair<KeyT, std::uint32_t>;
   thrust::device_vector<T> in(num_items);
 
   thrust::tabulate(in.begin(), in.end(), CountToPair<KeyT>{});
@@ -584,8 +584,9 @@ int main(int argc, char **argv)
   // Initialize device
   CubDebugExit(args.DeviceInit());
 
-  Test<std::uint8_t>();
-  Test<std::uint16_t>();
+  // NVBug 4136386
+  // Test<std::uint8_t>();
+  // Test<std::uint16_t>();
   Test<std::uint32_t>();
   Test<std::uint64_t>();