[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/loop_dispatch/dispatch_on_streams.cu

@@ -0,0 +1,45 @@
+//===----------------------------------------------------------------------===//
+//
+// 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/__places/place_partition.cuh>
+#include <cuda/experimental/__stf/internal/stf_places_partition_into_stf.cuh>
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+int main()
+{
+#if _CCCL_CTK_BELOW(12, 4)
+  fprintf(stderr, "Green contexts are not supported by this version of CUDA: skipping test.\n");
+  return 0;
+#else // ^^^ _CCCL_CTK_BELOW(12, 4) ^^^ / vvv _CCCL_CTK_AT_LEAST(12, 4) vvv
+
+  context ctx;
+
+  auto lX = ctx.logical_data<int>(size_t(32 * 1024 * 1024));
+
+  ctx.parallel_for(lX.shape(), lX.write())->*[] __device__(size_t i, auto x) {
+    x(i) = 3 * i - 7;
+  };
+
+  for (auto& sub_place :
+       place_partition(exec_place::current_device(), ctx.async_resources(), place_partition_scope::green_context))
+  {
+    for (size_t i = 0; i < 4; i++)
+    {
+      ctx.parallel_for(sub_place, lX.shape(), lX.rw())->*[] __device__(size_t i, auto x) {
+        x(i) += 1;
+      };
+    }
+  }
+
+  ctx.finalize();
+#endif // ^^^ _CCCL_CKT_AT_LEAST(12, 4) ^^^
+}

cccl_upstream/cudax/test/stf/loop_dispatch/loop_dispatch.cu

@@ -0,0 +1,41 @@
+//===----------------------------------------------------------------------===//
+//
+// 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/internal/loop_dispatch.cuh>
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+int main()
+{
+  context ctx;
+
+  // Loop count
+  int n = 1024;
+
+  auto lB = ctx.logical_data<int>(size_t(1024 * 1024));
+  ctx.parallel_for(lB.shape(), lB.write())->*[] __device__(size_t i, auto b) {
+    b(i) = 42;
+  };
+
+  loop_dispatch(ctx, exec_place::all_devices(), 0, n, [&](size_t) {
+    auto lA = ctx.logical_data<int>(size_t(1024 * 1024));
+
+    ctx.parallel_for(ctx.current_exec_place(), lA.shape(), lA.write())->*[] __device__(size_t i, auto a) {
+      a(i) = (int) (10.0 * cos((double) i));
+    };
+
+    ctx.parallel_for(ctx.current_exec_place(), lA.shape(), lA.rw(), lB.read())->*[] __device__(size_t i, auto a, auto b) {
+      a(i) += b(i);
+    };
+  });
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/loop_dispatch/nested_loop_dispatch.cu

@@ -0,0 +1,52 @@
+//===----------------------------------------------------------------------===//
+//
+// 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/std/cmath>
+
+#include <cuda/experimental/__stf/internal/loop_dispatch.cuh>
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+int main()
+{
+  context ctx;
+
+  auto lB = ctx.logical_data<int>(size_t(1024 * 1024));
+
+  ctx.parallel_for(lB.shape(), lB.write())->*[] __device__(size_t i, auto b) {
+    b(i) = 42;
+  };
+
+  // A fake grid which should work regardless of the underlying machine
+  auto grid = exec_place::repeat(exec_place::current_device(), 8);
+
+  // Split the affinity into 4 parts
+  loop_dispatch(ctx, grid, place_partition_scope::cuda_device, 0, 4, [&](size_t) {
+    // We should have 2 places per subplace
+    EXPECT(ctx.current_affinity().size() == 2);
+
+    // This should use ctx.current_affinity() implicitly
+    loop_dispatch(ctx, 0, 4, [&](size_t) {
+      auto lA = ctx.logical_data<int>(size_t(1024 * 1024));
+
+      ctx.parallel_for(ctx.current_exec_place(), lA.shape(), lA.write())->*[] __device__(size_t i, auto a) {
+        a(i) = (int) (10.0 * cuda::std::cos((double) i));
+      };
+
+      ctx.parallel_for(ctx.current_exec_place(), lA.shape(), lA.rw(), lB.read())
+          ->*[] __device__(size_t i, auto a, auto b) {
+                a(i) += b(i);
+              };
+    });
+  });
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/loop_dispatch/stackable_loop_dispatch.cu

@@ -0,0 +1,50 @@
+//===----------------------------------------------------------------------===//
+//
+// 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/internal/loop_dispatch.cuh>
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+int main()
+{
+  stackable_ctx ctx;
+  _CCCL_ASSERT(ctx.has_head_set(), "ctx construction must set head for current thread");
+
+  // Loop count
+  int n = 1024;
+
+  auto lB = ctx.logical_data<int>(size_t(1024 * 1024)).set_symbol("B");
+  ctx.parallel_for(lB.shape(), lB.write())->*[] __device__(size_t i, auto b) {
+    b(i) = 42;
+  };
+
+  lB.set_read_only(true);
+
+  for (size_t iter = 0; iter < 4; iter++)
+  {
+    loop_dispatch(ctx, exec_place::all_devices(), place_partition_scope::green_context, 0, n, [&](size_t iter) {
+      auto lA = ctx.logical_data<int>(size_t(1024 * 1024)).set_symbol(::std::string("A") + ::std::to_string(iter));
+
+      ctx.parallel_for(ctx.current_exec_place(), lA.shape(), lA.write()).set_symbol("pfor1" + ::std::to_string(iter))
+          ->*[] __device__(size_t i, auto a) {
+                a(i) = (int) (10.0 * cos((double) i));
+              };
+
+      ctx.parallel_for(ctx.current_exec_place(), lA.shape(), lA.rw(), lB.read())
+          .set_symbol("pfor2" + ::std::to_string(iter))
+          ->*[] __device__(size_t i, auto a, auto b) {
+                a(i) += b(i);
+              };
+    });
+  }
+
+  ctx.finalize();
+}