[INFRA] Import NVIDIA/CCCL upstream as optimization reference library

CCCL (CUDA C++ Core Libraries) provides:
- CUB: device/block/warp-level GPU primitives (reduce, scan, sort, topk)
- Thrust: high-level parallel algorithms (transform_reduce, sort, scan)
- libcudacxx: CUDA C++ standard library (atomics, barriers, memory)
- cudax: experimental features (memory resources, allocators)
- Tuning policies: per-SM hardware-specific algorithm parameters

Competition optimization vectors mapped to CCCL:
- Output TPS (83% weight): warp_reduce, block_reduce, device_topk
- Input TPS (14% weight): device_scan, block_load, prefetch
- Cache TPS (3% weight): prefix caching strategy patterns
- Memory (0.9 util): pooled/cached/buddy allocators

Source: https://github.com/NVIDIA/cccl (shallow clone, HEAD only)
License: Apache-2.0

cccl_upstream/cudax/test/stf/reclaiming/graph.cu

@@ -0,0 +1,74 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of CUDASTF in CUDA C++ Core Libraries,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+
+#include <iostream>
+
+#if !_CCCL_COMPILER(MSVC)
+using namespace cuda::experimental::stf;
+
+__global__ void kernel()
+{
+  // No-op
+}
+#endif // !_CCCL_COMPILER(MSVC)
+
+int main([[maybe_unused]] int argc, [[maybe_unused]] char** argv)
+{
+// TODO fix setenv
+#if !_CCCL_COMPILER(MSVC)
+  int nblocks       = 4;
+  size_t block_size = 1024 * 1024;
+
+  if (argc > 1)
+  {
+    nblocks = atoi(argv[1]);
+  }
+
+  if (argc > 2)
+  {
+    block_size = atoi(argv[2]);
+  }
+
+  // At most 1 buffer is allocated at the same time
+  setenv("MAX_ALLOC_CNT", "1", 0);
+
+  graph_ctx ctx;
+
+  ::std::vector<logical_data<slice<char>>> handles(nblocks);
+
+  char* h_buffer = new char[nblocks * block_size];
+
+  for (int i = 0; i < nblocks; i++)
+  {
+    handles[i] = ctx.logical_data(make_slice(&h_buffer[i * block_size], block_size));
+    handles[i].set_symbol("D_" + std::to_string(i));
+  }
+
+  // We only 2 buffers, we are forced to reuse the buffer from D0 for D2
+  for (int i = 0; i < 3; i++)
+  {
+    ctx.task(handles[i % nblocks].rw())->*[&](cudaStream_t s, auto /*unused*/) {
+      kernel<<<1, 1, 0, s>>>();
+    };
+  }
+
+  ctx.submit();
+
+  if (argc > 3)
+  {
+    std::cout << "Generating DOT output in " << argv[3] << '\n';
+    ctx.print_to_dot(argv[1]);
+  }
+
+  ctx.finalize();
+#endif // !_CCCL_COMPILER(MSVC)
+}

cccl_upstream/cudax/test/stf/reclaiming/graph_2.cu

@@ -0,0 +1,68 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of CUDASTF in CUDA C++ Core Libraries,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+
+#include <iostream>
+
+using namespace cuda::experimental::stf;
+
+__global__ void kernel()
+{
+  // No-op
+}
+
+int main(int argc, char** argv)
+{
+  int nblocks       = 4;
+  size_t block_size = 1024 * 1024;
+
+  if (argc > 1)
+  {
+    nblocks = atoi(argv[1]);
+  }
+
+  if (argc > 2)
+  {
+    block_size = atoi(argv[2]);
+  }
+
+  // At most 1 buffer is allocated at the same time
+  setenv("MAX_ALLOC_CNT", "2", 1);
+
+  graph_ctx ctx;
+  ::std::vector<logical_data<slice<char>>> handles(nblocks);
+  std::vector<char> h_buffer(nblocks * block_size);
+
+  for (int i = 0; i < nblocks; i++)
+  {
+    handles[i] = ctx.logical_data(make_slice(&h_buffer[i * block_size], block_size));
+    handles[i].set_symbol("D_" + std::to_string(i));
+  }
+
+  // We only 2 buffers, we are forced to reuse the buffer from D0 for D2
+  for (int i = 0; i < 3; i++)
+  {
+    ctx.task(handles[i % nblocks].read(), handles[(i + 1) % nblocks].rw())
+        ->*[&](cudaStream_t s, auto /*unused*/, auto /*unused*/) {
+              kernel<<<1, 1, 0, s>>>();
+            };
+  }
+
+  ctx.submit();
+
+  if (argc > 3)
+  {
+    std::cout << "Generating DOT output in " << argv[3] << '\n';
+    ctx.print_to_dot(argv[3]);
+  }
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/reclaiming/graph_real_oom.cu

@@ -0,0 +1,79 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of CUDASTF in CUDA C++ Core Libraries,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+
+#include <iostream>
+
+using namespace cuda::experimental::stf;
+
+__global__ void dummy() {}
+
+int main(int argc, char** argv)
+{
+  const int dev_id = 0;
+  cuda_safe_call(cudaSetDevice(dev_id));
+
+  cudaDeviceProp prop;
+  cuda_safe_call(cudaGetDeviceProperties(&prop, dev_id));
+
+  const size_t total_mem_ref = prop.totalGlobalMem;
+
+  size_t free_mem, total_mem;
+  cuda_safe_call(cudaMemGetInfo(&free_mem, &total_mem));
+
+  std::cout << "Device memory: " << total_mem_ref / (1024 * 1024.) << " MB"
+            << " FREE/TOTAL=" << free_mem / (1024 * 1024.) << "/" << total_mem / (1024 * 1024.) << '\n';
+
+  // Warning: this should represent 6% of device's available memory
+  size_t block_size = free_mem / 32;
+  int nblocks       = 2;
+
+  // We preallocate most available device memory to avoid being limited by host memory.
+  void* wasted_mem;
+  size_t wasted_size = block_size * 29;
+  cuda_safe_call(cudaMalloc(&wasted_mem, wasted_size));
+
+  std::cout << "Wasted: " << wasted_size / (1024 * 1024.) << " MB" << '\n';
+
+  graph_ctx ctx;
+
+  std::vector<logical_data<slice<char>>> handles(nblocks);
+
+  char* h_buffer = new char[nblocks * block_size];
+
+  for (int i = 0; i < 2; i++)
+  {
+    handles[i] = ctx.logical_data(make_slice(&h_buffer[i * block_size], block_size));
+    handles[i].set_symbol("D_" + std::to_string(i));
+  }
+
+  // There can only be a single handle on the device at the same time
+  for (int i = 0; i < nblocks; i++)
+  {
+    ctx.task(handles[i % nblocks].rw())->*[](cudaStream_t stream, auto /*unused*/) {
+      dummy<<<1, 1, 0, stream>>>();
+    };
+  }
+
+  // Checking the amount of memory actually available now
+  cuda_safe_call(cudaMemGetInfo(&free_mem, &total_mem));
+  std::cout << "Device memory: FREE/TOTAL=" << free_mem / (1024 * 1024.) << "/" << total_mem / (1024 * 1024.) << '\n';
+
+  ctx.submit();
+
+  if (argc > 1)
+  {
+    std::cout << "Generating DOT output in " << argv[1] << '\n';
+    ctx.print_to_dot(argv[1]);
+  }
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/reclaiming/stream.cu

@@ -0,0 +1,145 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of CUDASTF in CUDA C++ Core Libraries,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+__global__ void kernel(int i, slice<char> buf)
+{
+  buf[0] = (char) i;
+}
+
+static __global__ void cuda_sleep_kernel(long long int clock_cnt)
+{
+  long long int start_clock  = clock64();
+  long long int clock_offset = 0;
+  while (clock_offset < clock_cnt)
+  {
+    clock_offset = clock64() - start_clock;
+  }
+}
+
+void cuda_sleep(double ms, cudaStream_t stream)
+{
+  int device;
+  cudaGetDevice(&device);
+
+  // cudaDevAttrClockRate: Peak clock frequency in kilohertz;
+  int clock_rate;
+  cudaDeviceGetAttribute(&clock_rate, cudaDevAttrClockRate, device);
+
+  long long int clock_cnt = (long long int) (ms * clock_rate);
+  cuda_sleep_kernel<<<1, 1, 0, stream>>>(clock_cnt);
+}
+
+// Dummy allocator which only allocates a single block on a single device
+//
+// The first allocation succeeds, next attempt will fail until buffer is
+// deallocated.
+class one_block_allocator : public block_allocator_interface
+{
+public:
+  one_block_allocator() = default;
+
+public:
+  // Note that this allocates memory immediately, so we just do not modify the event list and ignore it
+  void* allocate(backend_ctx_untyped&, const data_place& memory_node, ::std::ptrdiff_t& s, event_list&) override
+  {
+    if (busy)
+    {
+      s = -s;
+      return nullptr;
+    }
+
+    EXPECT(memory_node.is_device());
+    if (!base)
+    {
+      cuda_safe_call(cudaMalloc(&base, s));
+    }
+
+    busy = true;
+
+    return base;
+  }
+
+  void deallocate(backend_ctx_untyped&, const data_place&, event_list&, void*, size_t) override
+  {
+    EXPECT(busy);
+    busy = false;
+  }
+
+  event_list deinit(backend_ctx_untyped&) override
+  {
+    return event_list();
+  }
+
+  std::string to_string() const override
+  {
+    return "dummy";
+  }
+
+private:
+  // We have a single block, so we keep its address, and a flag to indicate
+  // if it's busy
+  void* base = nullptr;
+  bool busy  = false;
+};
+
+int main(int argc, char** argv)
+{
+  int nblocks       = 4;
+  size_t block_size = 1024 * 1024;
+
+  if (argc > 1)
+  {
+    nblocks = atoi(argv[1]);
+  }
+
+  if (argc > 2)
+  {
+    block_size = atoi(argv[2]);
+  }
+
+  stream_ctx ctx;
+
+  auto dummy_alloc = block_allocator<one_block_allocator>(ctx);
+  ctx.set_allocator(dummy_alloc);
+
+  ::std::vector<logical_data<slice<char>>> handles(nblocks);
+
+  EXPECT(nblocks > 0);
+  EXPECT(block_size > 0);
+
+  char* h_buffer = new char[nblocks * block_size];
+
+  for (int i = 0; i < nblocks; i++)
+  {
+    handles[i] = ctx.logical_data(make_slice(&h_buffer[i * block_size], block_size));
+    handles[i].set_symbol("D_" + std::to_string(i));
+  }
+
+  // We only 2 buffers, we are forced to reuse the buffer from D0 for D2
+  for (int i = 0; i < nblocks; i++)
+  {
+    ctx.task(handles[i % nblocks].rw())->*[&](cudaStream_t s, auto buf) {
+      // Wait 100ms to have a more stressful asynchronous execution
+      cuda_sleep(100, s);
+      kernel<<<1, 1, 0, s>>>(i, buf);
+    };
+  }
+
+  ctx.finalize();
+
+  for (int i = 0; i < nblocks; i++)
+  {
+    EXPECT(h_buffer[block_size * i] == i);
+  }
+}