[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/algorithm/algorithm_with_read.cu

@@ -0,0 +1,58 @@
+//===----------------------------------------------------------------------===//
+//
+// 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;
+
+template <typename T>
+void init(context& ctx, logical_data<T> l, int val)
+{
+  ctx.parallel_for(l.shape(), l.write())->*[=] __device__(size_t i, auto s) {
+    s(i) = val;
+  };
+}
+
+int main()
+{
+  context ctx;
+
+  auto a = ctx.logical_data<int>(10000000);
+  auto b = ctx.logical_data<int>(10000000);
+  auto c = ctx.logical_data<int>(10000000);
+  auto d = ctx.logical_data<int>(10000000);
+
+  init(ctx, a, 12);
+  init(ctx, b, 35);
+  init(ctx, c, 42);
+  init(ctx, d, 17);
+
+  /* a += 1; a += b; */
+  auto fn = [](context ctx, logical_data<slice<int>> a, logical_data<slice<int>> b) {
+    ctx.parallel_for(a.shape(), a.rw())->*[] __device__(size_t i, auto sa) {
+      sa(i) += 1;
+    };
+    ctx.parallel_for(a.shape(), a.rw(), b.read())->*[] __device__(size_t i, auto sa, auto sb) {
+      sa(i) += sb(i);
+    };
+  };
+
+  algorithm alg;
+
+  for (size_t i = 0; i < 100; i++)
+  {
+    alg.run_as_task(fn, ctx, a.rw(), b.read());
+    alg.run_as_task(fn, ctx, a.rw(), c.read());
+    alg.run_as_task(fn, ctx, c.rw(), d.read());
+    alg.run_as_task(fn, ctx, d.rw(), a.read());
+  }
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/algorithm/graph_algorithms.cu

@@ -0,0 +1,108 @@
+//===----------------------------------------------------------------------===//
+//
+// 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;
+
+template <typename ctx_t, typename T, typename T2>
+void lib_call(ctx_t& ctx, logical_data<T> a, logical_data<T2> b)
+{
+  nvtx_range r("lib_call");
+  // b *= 2
+  // a = a + b
+  // b = a
+  ctx.parallel_for(b.shape(), b.rw())->*[] __device__(size_t i, auto sb) {
+    sb(i) *= 2;
+  };
+
+  ctx.parallel_for(a.shape(), a.rw(), b.read())->*[] __device__(size_t i, auto sa, auto sb) {
+    sa(i) += sb(i);
+  };
+
+  ctx.parallel_for(a.shape(), a.read(), b.write())->*[] __device__(size_t i, auto sa, auto sb) {
+    sb(i) = sa(i);
+  };
+}
+
+template <typename T>
+void init(context& ctx, logical_data<T> l, int val)
+{
+  ctx.parallel_for(l.shape(), l.write())->*[=] __device__(size_t i, auto s) {
+    s(i) = val;
+  };
+}
+
+int main()
+{
+  context ctx;
+
+  auto a = ctx.logical_data<int>(size_t(10000000));
+  auto b = ctx.logical_data<int>(size_t(10000000));
+  auto c = ctx.logical_data<int>(size_t(10000000));
+  auto d = ctx.logical_data<int>(size_t(10000000));
+
+  init(ctx, a, 12);
+  init(ctx, b, 35);
+  init(ctx, c, 42);
+  init(ctx, d, 42);
+
+  auto fn = [](context ctx, logical_data<slice<int>> a, logical_data<slice<int>> b) {
+    ctx.parallel_for(b.shape(), b.rw())->*[] __device__(size_t i, auto sb) {
+      sb(i) *= 2;
+    };
+    ctx.parallel_for(a.shape(), a.rw(), b.read())->*[] __device__(size_t i, auto sa, auto sb) {
+      sa(i) += sb(i);
+    };
+    ctx.parallel_for(a.shape(), a.read(), b.write())->*[] __device__(size_t i, auto sa, auto sb) {
+      sb(i) = sa(i);
+    };
+  };
+
+  algorithm alg;
+
+  {
+    nvtx_range r("run");
+    for (size_t i = 0; i < 100; i++)
+    {
+      ctx.task(a.rw(), b.rw())->*[&alg, &fn, &ctx](cudaStream_t stream, slice<int> sa, slice<int> sb) {
+        alg.run(fn, ctx, stream, sa, sb);
+      };
+
+      ctx.task(b.rw(), c.rw())->*[&alg, &fn, &ctx](cudaStream_t stream, slice<int> sb, slice<int> sc) {
+        alg.run(fn, ctx, stream, sb, sc);
+      };
+
+      ctx.task(c.rw(), d.rw())->*[&alg, &fn, &ctx](cudaStream_t stream, slice<int> sc, slice<int> sd) {
+        alg.run(fn, ctx, stream, sc, sd);
+      };
+
+      ctx.task(d.rw(), a.rw())->*[&alg, &fn, &ctx](cudaStream_t stream, slice<int> sd, slice<int> sa) {
+        alg.run(fn, ctx, stream, sd, sa);
+      };
+    }
+  }
+
+  {
+    nvtx_range r("run_as_task");
+    for (size_t i = 0; i < 100; i++)
+    {
+      alg.run_as_task(fn, ctx, a.rw(), b.rw());
+
+      alg.run_as_task(fn, ctx, a.rw(), c.rw());
+
+      alg.run_as_task(fn, ctx, c.rw(), d.rw());
+
+      alg.run_as_task(fn, ctx, d.rw(), a.rw());
+    }
+  }
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/algorithm/in_graph_ctx.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.cuh>
+
+using namespace cuda::experimental::stf;
+
+template <typename context_t, typename T>
+void init(context_t& ctx, logical_data<T> l, int val)
+{
+  ctx.parallel_for(l.shape(), l.write())->*[=] __device__(size_t i, auto s) {
+    s(i) = val;
+  };
+}
+
+int main()
+{
+  //    context ctx = graph_ctx();
+  graph_ctx ctx;
+
+  auto a = ctx.logical_data<int>(size_t(1000000));
+  auto b = ctx.logical_data<int>(size_t(1000000));
+  auto c = ctx.logical_data<int>(size_t(1000000));
+  auto d = ctx.logical_data<int>(size_t(1000000));
+
+  init(ctx, a, 12);
+  init(ctx, b, 35);
+  init(ctx, c, 42);
+  init(ctx, d, 42);
+
+  auto fn = [](context ctx, logical_data<slice<int>> a, logical_data<slice<int>> b) {
+    ctx.parallel_for(a.shape(), a.rw())->*[] __device__(size_t i, auto sa) {
+      sa(i) *= 3;
+    };
+    ctx.parallel_for(b.shape(), b.rw())->*[] __device__(size_t i, auto sb) {
+      sb(i) *= 2;
+    };
+
+    ctx.parallel_for(a.shape(), a.rw(), b.read())->*[] __device__(size_t i, auto sa, auto sb) {
+      sa(i) += sb(i);
+    };
+
+    ctx.parallel_for(a.shape(), a.read(), b.write())->*[] __device__(size_t i, auto sa, auto sb) {
+      sb(i) = sa(i);
+    };
+  };
+
+  algorithm alg;
+
+  for (size_t i = 0; i < 5; i++)
+  {
+    alg.run_as_task(fn, ctx, a.rw(), b.rw());
+    alg.run_as_task(fn, ctx, a.rw(), c.rw());
+    alg.run_as_task(fn, ctx, c.rw(), d.rw());
+    alg.run_as_task(fn, ctx, d.rw(), a.rw());
+  }
+
+  ctx.finalize();
+
+  // cudaGraphDebugDotPrint(ctx.get_graph(), "pif.dot", 0);
+}

cccl_upstream/cudax/test/stf/algorithm/nested.cu

@@ -0,0 +1,67 @@
+//===----------------------------------------------------------------------===//
+//
+// 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;
+
+template <typename T>
+void init(context& ctx, logical_data<T> l, int val)
+{
+  ctx.parallel_for(l.shape(), l.write())->*[=] __device__(size_t i, auto s) {
+    s(i) = val;
+  };
+}
+
+int main()
+{
+  context ctx;
+
+  auto a = ctx.logical_data<int>(size_t(1000000));
+  auto b = ctx.logical_data<int>(size_t(1000000));
+  auto c = ctx.logical_data<int>(size_t(1000000));
+  auto d = ctx.logical_data<int>(size_t(1000000));
+
+  init(ctx, a, 12);
+  init(ctx, b, 35);
+  init(ctx, c, 42);
+  init(ctx, d, 17);
+
+  auto fn1 = [](context ctx, logical_data<slice<int>> a) {
+    ctx.parallel_for(a.shape(), a.rw())->*[] __device__(size_t i, auto sa) {
+      sa(i) += 1;
+    };
+  };
+
+  algorithm alg1;
+
+  auto fn2 = [&alg1, &fn1](context ctx, logical_data<slice<int>> a, logical_data<slice<int>> b) {
+    alg1.run_as_task(fn1, ctx, a.rw());
+    alg1.run_as_task(fn1, ctx, b.rw());
+    ctx.parallel_for(a.shape(), a.rw(), b.read())->*[] __device__(size_t i, auto sa, auto sb) {
+      sa(i) += sb(i);
+    };
+    ctx.parallel_for(a.shape(), a.read(), b.write())->*[] __device__(size_t i, auto sa, auto sb) {
+      sb(i) = sa(i);
+    };
+  };
+
+  algorithm alg2;
+
+  for (size_t i = 0; i < 100; i++)
+  {
+    alg2.run_as_task(fn2, ctx, a.rw(), b.rw());
+    alg2.run_as_task(fn2, ctx, a.rw(), c.rw());
+    alg2.run_as_task(fn2, ctx, c.rw(), d.rw());
+    alg2.run_as_task(fn2, ctx, d.rw(), a.rw());
+  }
+
+  ctx.finalize();
+}