[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/reductions/many_inc.cu

@@ -0,0 +1,51 @@
+//===----------------------------------------------------------------------===//
+//
+// 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/stream/interfaces/slice_reduction_ops.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+__global__ void add(int* ptr)
+{
+  *ptr = *ptr + 1;
+}
+
+int main()
+{
+  stream_ctx ctx;
+
+  auto redux_op = std::make_shared<slice_reduction_op_sum<int>>();
+
+  int a = 17;
+
+  auto handle = ctx.logical_data(make_slice(&a, 1));
+
+  int ndevs;
+  cuda_safe_call(cudaGetDeviceCount(&ndevs));
+
+  int K = 1024;
+  for (int i = 0; i < K; i++)
+  {
+    // Increment the variable by 1
+    ctx.task(exec_place::device(i % ndevs), handle.relaxed(redux_op))->*[](auto stream, auto s) {
+      add<<<1, 1, 0, stream>>>(s.data_handle());
+    };
+  }
+
+  // Total value should be initial value + K
+  ctx.task(exec_place::host(), handle.read())->*[&](auto stream, auto s) {
+    cuda_safe_call(cudaStreamSynchronize(stream));
+    EXPECT(s(0) == 17 + K);
+    // printf("VALUE %d expected %d\n", s(0), 17 + K);
+  };
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/reductions/reduce_sum.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.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ *
+ * @brief Test reduce access mode
+ *
+ */
+
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+template <typename context_t>
+void run()
+{
+  context_t ctx;
+  auto lsum = ctx.logical_data(shape_of<scalar_view<size_t>>());
+
+  size_t N = 100000;
+
+  ctx.parallel_for(box(N), lsum.reduce(reducer::sum<size_t>{}))->*[] __device__(size_t i, auto& sum) {
+    sum++;
+  };
+
+  size_t res_sum = ctx.wait(lsum);
+
+  ctx.finalize();
+
+  _CCCL_ASSERT(res_sum == N, "Invalid result");
+}
+
+int main()
+{
+  run<stream_ctx>();
+  run<graph_ctx>();
+}

cccl_upstream/cudax/test/stf/reductions/redux_test.cu

@@ -0,0 +1,42 @@
+//===----------------------------------------------------------------------===//
+//
+// 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/stream/interfaces/slice_reduction_ops.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+__global__ void add(int* ptr, int value)
+{
+  *ptr = *ptr + value;
+}
+
+int main()
+{
+  stream_ctx ctx;
+
+  int a         = 17;
+  auto handle   = ctx.logical_data(make_slice(&a, 1));
+  auto redux_op = std::make_shared<slice_reduction_op_sum<int>>();
+  ctx.task(handle.relaxed(redux_op))->*[](auto stream, auto s) {
+    add<<<1, 1, 0, stream>>>(s.data_handle(), 42);
+  };
+
+  ctx.task(handle.rw())->*[](auto stream, auto s) {
+    add<<<1, 1, 0, stream>>>(s.data_handle(), 1);
+  };
+
+  ctx.task(exec_place::host(), handle.read())->*[](auto stream, auto s) {
+    cuda_safe_call(cudaStreamSynchronize(stream));
+    EXPECT(s(0) == 17 + 42 + 1);
+  };
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/reductions/redux_test2.cu

@@ -0,0 +1,78 @@
+//===----------------------------------------------------------------------===//
+//
+// 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/stream/interfaces/slice_reduction_ops.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+__global__ void add(int* ptr, int value)
+{
+  *ptr = *ptr + value;
+}
+
+__global__ void check_val(const int* ptr, int expected)
+{
+  assert(*ptr == expected);
+}
+
+/*
+ * This test ensures that we can reconstruct a piece of data where there are
+ * multiple "shared" instances, and "redux" instances
+ */
+int main()
+{
+  stream_ctx ctx;
+
+  int ndevs;
+  cuda_safe_call(cudaGetDeviceCount(&ndevs));
+
+  if (ndevs < 2)
+  {
+    fprintf(stderr, "Skipping test: need at least 2 devices.\n");
+    return 0;
+  }
+
+  auto redux_op = std::make_shared<slice_reduction_op_sum<int>>();
+
+  int a = 17;
+
+  // init op (17)
+  auto handle = ctx.logical_data(make_slice(&a, 1));
+
+  // RW dev0 (18)
+  ctx.task(exec_place::device(0), handle.rw())->*[](auto stream, auto s) {
+    add<<<1, 1, 0, stream>>>(s.data_handle(), 1);
+  };
+
+  // READ dev1 (18)
+  ctx.task(exec_place::device(1), handle.read())->*[](auto stream, auto s) {
+    check_val<<<1, 1, 0, stream>>>(s.data_handle(), 18);
+  };
+
+  // REDUX dev1 (18 + 42)
+  ctx.task(exec_place::device(1), handle.relaxed(redux_op))->*[](auto stream, auto s) {
+    add<<<1, 1, 0, stream>>>(s.data_handle(), 42);
+  };
+
+  // READ dev0 (18 + 42)
+  ctx.task(exec_place::device(0), handle.read())->*[](auto stream, auto s) {
+    check_val<<<1, 1, 0, stream>>>(s.data_handle(), 18 + 42);
+  };
+
+  // READ
+  ctx.task(exec_place::host(), handle.read())->*[](auto stream, auto s) {
+    cuda_safe_call(cudaStreamSynchronize(stream));
+    EXPECT(s(0) == 18 + 42);
+    // printf("VALUE %d expected %d\n", s(0), 18 + 42);
+  };
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/reductions/slice2d_reduction.cu

@@ -0,0 +1,59 @@
+//===----------------------------------------------------------------------===//
+//
+// 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/stream/interfaces/slice_reduction_ops.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+__global__ void add(slice<int, 2> s, int val)
+{
+  size_t tid      = threadIdx.x + blockIdx.x * blockDim.x;
+  size_t nthreads = blockDim.x * gridDim.x;
+
+  for (size_t j = 0; j < s.extent(1); j++)
+  {
+    for (size_t i = tid; i < s.extent(0); i += nthreads)
+    {
+      s(i, j) += val;
+    }
+  }
+}
+
+int main()
+{
+  stream_ctx ctx;
+
+  int array[6] = {0, 1, 2, 3, 4, 5};
+
+  // auto handle = ctx.logical_data(slice<int, 2>(&array[0], std::tuple{ 2, 3 }, 2));
+  auto handle = ctx.logical_data(make_slice(&array[0], std::tuple{2, 3}, 2));
+
+  auto redux_op = std::make_shared<slice_reduction_op_sum<int, 2>>();
+
+  ctx.task(handle.relaxed(redux_op))->*[](auto stream, auto s) {
+    add<<<32, 32, 0, stream>>>(s, 42);
+  };
+
+  ctx.task(exec_place::host(), handle.read())->*[](auto stream, auto s) {
+    cuda_safe_call(cudaStreamSynchronize(stream));
+
+    for (size_t j = 0; j < s.extent(1); j++)
+    {
+      for (size_t i = 0; i < s.extent(0); i++)
+      {
+        // fprintf(stderr, "%d\t", s(i, j));
+      }
+      // fprintf(stderr, "\n");
+    }
+  };
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/reductions/slice_custom_op.cu

@@ -0,0 +1,72 @@
+//===----------------------------------------------------------------------===//
+//
+// 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/stream/interfaces/slice_reduction_ops.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+struct OR_op
+{
+  static void init_host(bool& out)
+  {
+    out = false;
+  };
+  static __device__ void init_gpu(bool& out)
+  {
+    out = false;
+  };
+  static void op_host(const bool& in, bool& inout)
+  {
+    inout |= in;
+  };
+  static __device__ void op_gpu(const bool& in, bool& inout)
+  {
+    inout |= in;
+  };
+};
+
+int main()
+{
+  stream_ctx ctx;
+
+  int ndevs;
+  cuda_safe_call(cudaGetDeviceCount(&ndevs));
+
+  bool A[4] = {false, false, true, false};
+  bool B[4] = {false, false, false, false};
+  bool C[4] = {true, false, false, true};
+
+  auto lA = ctx.logical_data(make_slice(&A[0], 4));
+  auto lB = ctx.logical_data(make_slice(&B[0], 4));
+  auto lC = ctx.logical_data(make_slice(&C[0], 4));
+
+  auto op = std::make_shared<slice_reduction_op<bool, 1, OR_op>>();
+
+  // C |= A
+  ctx.task(lC.relaxed(op), lA.read())->*[](auto stream, auto sC, auto sA) {
+    cudaMemcpyAsync(sC.data_handle(), sA.data_handle(), sA.extent(0) * sizeof(bool), cudaMemcpyDeviceToDevice, stream);
+  };
+
+  // C |= B
+  ctx.task(lC.relaxed(op), lB.read())->*[](auto stream, auto sC, auto sB) {
+    cudaMemcpyAsync(sC.data_handle(), sB.data_handle(), sB.extent(0) * sizeof(bool), cudaMemcpyDeviceToDevice, stream);
+  };
+
+  ctx.task(exec_place::host(), lC.read())->*[](auto stream, auto sC) {
+    cuda_safe_call(cudaStreamSynchronize(stream));
+    for (size_t i = 0; i < sC.extent(0); i++)
+    {
+      // fprintf(stderr, "RESULT C[i] = %d\n", sC(i));
+    }
+  };
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/reductions/successive_reductions.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/stream/interfaces/slice_reduction_ops.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+using scalar_t = slice_stream_interface<int, 1>;
+
+template <typename T>
+__global__ void set_value(T* addr, T val)
+{
+  *addr = val;
+}
+
+template <typename T>
+__global__ void check_value_and_reset(T* addr, T val, T reset_val)
+{
+  assert(*addr == val);
+  *addr = reset_val;
+}
+
+template <typename T>
+__global__ void add_val(T* inout_addr, T val)
+{
+  *inout_addr += val;
+}
+
+int main()
+{
+  stream_ctx ctx;
+  const int N = 4;
+
+  auto var_handle = ctx.logical_data(shape_of<slice<int>>(1));
+  var_handle.set_symbol("var");
+
+  auto redux_op = std::make_shared<slice_reduction_op_sum<int>>();
+
+  const int niters = 4;
+  for (int iter = 0; iter < niters; iter++)
+  {
+    // We add i (total = N(N-1)/2 + initial_value)
+    for (int i = 0; i < N; i++)
+    {
+      ctx.task(var_handle.relaxed(redux_op))->*[=](cudaStream_t stream, auto d_var) {
+        add_val<<<1, 1, 0, stream>>>(d_var.data_handle(), i);
+        cuda_safe_call(cudaGetLastError());
+      };
+    }
+
+    // Check that we have the expected value, and reset it so that we can perform another reduction
+    ctx.task(var_handle.rw())->*[=](cudaStream_t stream, auto d_var) {
+      int expected = (N * (N - 1)) / 2;
+      check_value_and_reset<<<1, 1, 0, stream>>>(d_var.data_handle(), expected, 0);
+    };
+  }
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/reductions/successive_reductions_pfor.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/__stf/stream/interfaces/slice_reduction_ops.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+int main()
+{
+  stream_ctx ctx;
+  const int N = 4;
+
+  auto var_handle = ctx.logical_data(shape_of<slice<int>>(1));
+  var_handle.set_symbol("var");
+
+  auto op = std::make_shared<slice_reduction_op_sum<int>>();
+
+  const int niters = 4;
+  for (int iter = 0; iter < niters; iter++)
+  {
+    // We add i (total = N(N-1)/2 + initial_value)
+    for (int i = 0; i < N; i++)
+    {
+      ctx.parallel_for(var_handle.shape(), var_handle.relaxed(op))->*[=] _CCCL_DEVICE(size_t ind, auto d_var) {
+        atomicAdd(d_var.data_handle(), i);
+      };
+    }
+
+    // Check that we have the expected value, and reset it so that we can perform another reduction
+    ctx.parallel_for(var_handle.shape(), var_handle.rw())->*[=] _CCCL_DEVICE(size_t ind, auto d_var) {
+      assert(d_var(0) == (N * (N - 1)) / 2);
+      d_var(0) = 0;
+    };
+  }
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/reductions/sum.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.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief Use the reduction access mode to add variables concurrently on different places
+ */
+
+#include <cuda/experimental/__stf/stream/reduction.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+using scalar_t = slice_stream_interface<int, 1>;
+
+template <typename T>
+__global__ void set_value(T* addr, T val)
+{
+  *addr = val;
+}
+
+template <typename T>
+__global__ void add(const T* in_addr, T* inout_addr)
+{
+  // printf("ADD: %d += %d : RES %d\n", *inout_addr, *in_addr, *inout_addr + *in_addr);
+
+  *inout_addr += *in_addr;
+}
+
+template <typename T>
+__global__ void add_val(T* inout_addr, T val)
+{
+  *inout_addr += val;
+}
+
+/*
+ * Define a SUM reduction operator over a scalar
+ */
+class scalar_sum_t : public stream_reduction_operator_untyped
+{
+public:
+  scalar_sum_t()
+      : stream_reduction_operator_untyped() {};
+
+  void stream_redux_op(
+    logical_data_untyped& d,
+    const data_place& /*unused*/,
+    instance_id_t inout_instance_id,
+    const data_place& /*unused*/,
+    instance_id_t in_instance_id,
+    const exec_place& /*unused*/,
+    cudaStream_t s) override
+  {
+    auto& in_instance    = d.instance<typename scalar_t::element_type>(in_instance_id);
+    auto& inout_instance = d.instance<typename scalar_t::element_type>(inout_instance_id);
+    // fprintf(stderr, "REDUX OP d %p inout (node %d id %d addr %p) in (node %d id %d addr %p)\n", d,
+    //        inout_memory_node, inout_instance_id, *inout_instance, in_memory_node, in_instance_id, *in_instance);
+    add<<<1, 1, 0, s>>>(in_instance.data_handle(), inout_instance.data_handle());
+  }
+
+  void stream_init_op(logical_data_untyped& d,
+                      const data_place& /*unused*/,
+                      instance_id_t out_instance_id,
+                      const exec_place& /*unused*/,
+                      cudaStream_t s) override
+  {
+    auto& out_instance = d.instance<typename scalar_t::element_type>(out_instance_id);
+    // fprintf(stderr, "REDUX INIT d %p memory node %d instance id %d => addr %p\n", d, out_memory_node,
+    //        out_instance_id, *out_instance);
+    set_value<<<1, 1, 0, s>>>(out_instance.data_handle(), 0);
+  }
+};
+
+int main()
+{
+  const int N = 4;
+
+  stream_ctx ctx;
+
+  auto var_handle = ctx.logical_data(shape_of<slice<int>>(1));
+  var_handle.set_symbol("var");
+
+  auto redux_op = std::make_shared<scalar_sum_t>();
+
+  // We add i (total = N(N-1)/2 + initial_value)
+  for (int i = 0; i < N; i++)
+  {
+    ctx.task(var_handle.relaxed(redux_op))->*[&](cudaStream_t stream, auto d_var) {
+      add_val<<<1, 1, 0, stream>>>(d_var.data_handle(), i);
+    };
+  }
+
+  // Check result
+  ctx.task(exec_place::host(), var_handle.read())->*[&](cudaStream_t stream, auto h_var) {
+    cuda_safe_call(cudaStreamSynchronize(stream));
+    int expected = (N * (N - 1)) / 2;
+    EXPECT(h_var(0) == expected);
+  };
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/reductions/sum_array.cu

@@ -0,0 +1,120 @@
+//===----------------------------------------------------------------------===//
+//
+// 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/stream/reduction.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+#include <iostream>
+
+using namespace cuda::experimental::stf;
+
+using scalar_t = slice_stream_interface<int, 1>;
+
+template <typename T>
+__global__ void set_value(T* addr, T val)
+{
+  *addr = val;
+}
+
+template <typename T>
+__global__ void add(const T* in_addr, T* inout_addr)
+{
+  *inout_addr += *in_addr;
+}
+
+/*
+ * Define a SUM reduction operator over a scalar
+ */
+class scalar_sum_t : public stream_reduction_operator_untyped
+{
+public:
+  scalar_sum_t()
+      : stream_reduction_operator_untyped() {};
+
+  void stream_redux_op(
+    logical_data_untyped& d,
+    const data_place& /*unused*/,
+    instance_id_t inout_instance_id,
+    const data_place& /*unused*/,
+    instance_id_t in_instance_id,
+    const exec_place& /*unused*/,
+    cudaStream_t s) override
+  {
+    auto& in_instance    = d.instance<typename scalar_t::element_type>(in_instance_id);
+    auto& inout_instance = d.instance<typename scalar_t::element_type>(inout_instance_id);
+    add<<<1, 1, 0, s>>>(in_instance.data_handle(), inout_instance.data_handle());
+  }
+
+  void stream_init_op(logical_data_untyped& d,
+                      const data_place& /*unused*/,
+                      instance_id_t out_instance_id,
+                      const exec_place& /*unused*/,
+                      cudaStream_t s) override
+  {
+    auto& out_instance = d.instance<typename scalar_t::element_type>(out_instance_id);
+    // fprintf(stderr, "REDUX INIT d %p memory node %d instance id %d => addr %p\n", d, out_memory_node,
+    //        out_instance_id, *out_instance);
+    set_value<<<1, 1, 0, s>>>(out_instance.data_handle(), 0);
+  }
+};
+
+int main()
+{
+  stream_ctx ctx;
+
+  const int N = 128;
+
+  // We have an array, and a handle for each entry of the array
+  int array[N];
+  logical_data<slice<int>> array_handles[N];
+
+  /*
+   * We are going to compute the sum of this array
+   */
+  for (int i = 0; i < N; i++)
+  {
+    array[i]         = i;
+    array_handles[i] = ctx.logical_data(&array[i], {1});
+    array_handles[i].set_symbol(std::string("array[") + std::to_string(i) + std::string("]"));
+  }
+
+  logical_data<slice<int>> var_handle = ctx.logical_data(shape_of<slice<int>>(1));
+  var_handle.set_symbol("var");
+
+  int check_sum = 0;
+  for (int i = 0; i < N; i++)
+  {
+    check_sum += array[i];
+  }
+
+  auto redux_op = std::make_shared<scalar_sum_t>();
+
+  for (int i = 0; i < N; i++)
+  {
+    ctx.task(var_handle.relaxed(redux_op), array_handles[i].read())
+        ->*[](cudaStream_t stream, auto d_var, auto d_array_i) {
+              add<<<1, 1, 0, stream>>>(d_array_i.data_handle(), d_var.data_handle());
+            };
+  }
+
+  // Force the reconstruction of data on the device, so that no transfers are
+  // necessary while reconstructing the result.
+  // This will of course not be necessary in the future ...
+  ctx.task(var_handle.read())->*[](cudaStream_t /*unused*/, auto /*unused*/) {};
+
+  // Check result
+  ctx.task(exec_place::host(), var_handle.read())->*[=](cudaStream_t stream, auto h_var) {
+    cuda_safe_call(cudaStreamSynchronize(stream));
+    int value = h_var(0);
+    EXPECT(value == check_sum);
+  };
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/reductions/sum_multiple_places.cu

@@ -0,0 +1,62 @@
+//===----------------------------------------------------------------------===//
+//
+// 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/stream/interfaces/slice_reduction_ops.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+template <typename T>
+__global__ void add_val(slice<T> inout, T val)
+{
+  inout(0) += val;
+}
+
+int main()
+{
+  stream_ctx ctx;
+
+  int ndevs;
+  cuda_safe_call(cudaGetDeviceCount(&ndevs));
+
+  const int N = 4;
+
+  int var = 42;
+
+  auto var_handle = ctx.logical_data(make_slice(&var, 1));
+  var_handle.set_symbol("var");
+
+  auto redux_op = std::make_shared<slice_reduction_op_sum<int>>();
+
+  // We add i twice (total = N(N-1) + initial_value)
+  for (int i = 0; i < N; i++)
+  {
+    // device
+    for (int d = 0; d < ndevs; d++)
+    {
+      ctx.task(exec_place::device(d), var_handle.relaxed(redux_op))->*[=](cudaStream_t s, auto var) {
+        add_val<int><<<1, 1, 0, s>>>(var, i);
+      };
+    }
+
+    // host
+    ctx.host_launch(var_handle.relaxed(redux_op))->*[=](auto var) {
+      var(0) += i;
+    };
+  }
+
+  // Check result
+  ctx.host_launch(var_handle.read())->*[&](auto var) {
+    int expected = 42 + (N * (N - 1)) / 2 * (ndevs + 1);
+    EXPECT(var(0) == expected);
+  };
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/reductions/sum_multiple_places_no_refvalue.cu

@@ -0,0 +1,118 @@
+//===----------------------------------------------------------------------===//
+//
+// 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.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief This test ensures that we can apply a reduction on a logical data
+ *        described using a shape
+ */
+
+#include <cuda/experimental/__stf/stream/reduction.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+using scalar_t = slice<int>;
+
+template <typename T>
+__global__ void set_value(T* addr, T val)
+{
+  *addr = val;
+}
+
+// instance id are put for debugging purpose ...
+template <typename T>
+__global__ void add(const T* in_addr, T* inout_addr)
+{
+  *inout_addr += *in_addr;
+}
+
+template <typename T>
+__global__ void add_val(T* inout_addr, T val)
+{
+  *inout_addr += val;
+}
+
+class scalar_sum_t : public stream_reduction_operator<scalar_t>
+{
+public:
+  void op(const scalar_t& in, scalar_t& inout, const exec_place& e, cudaStream_t s) override
+  {
+    if (e.affine_data_place().is_host())
+    {
+      // TODO make a callback when the situation gets better
+      cuda_safe_call(cudaStreamSynchronize(s));
+      *inout.data_handle() += *in.data_handle();
+    }
+    else
+    {
+      // this is not the host, so this has to be a device ... (XXX)
+      add<int><<<1, 1, 0, s>>>(in.data_handle(), inout.data_handle());
+    }
+  }
+
+  void init_op(scalar_t& out, const exec_place& e, cudaStream_t s) override
+  {
+    if (e.affine_data_place().is_host())
+    {
+      // TODO make a callback when the situation gets better
+      cuda_safe_call(cudaStreamSynchronize(s));
+      *out.data_handle() = 0;
+    }
+    else
+    {
+      // this is not the host, so this has to be a device ... (XXX)
+      set_value<int><<<1, 1, 0, s>>>(out.data_handle(), 0);
+    }
+  }
+};
+
+int main()
+{
+  stream_ctx ctx;
+
+  int ndevs;
+  cuda_safe_call(cudaGetDeviceCount(&ndevs));
+
+  const int N = 4;
+
+  auto var_handle = ctx.logical_data(shape_of<slice<int>>(1));
+  var_handle.set_symbol("var");
+
+  auto redux_op = std::make_shared<scalar_sum_t>();
+
+  // We add i twice (total = N(N-1) + initial_value)
+  for (int i = 0; i < N; i++)
+  {
+    // device
+    for (int d = 0; d < ndevs; d++)
+    {
+      ctx.task(exec_place::device(d), var_handle.relaxed(redux_op))->*[&](cudaStream_t s, auto var) {
+        add_val<int><<<1, 1, 0, s>>>(var.data_handle(), i);
+      };
+    }
+
+    // host
+    ctx.task(exec_place::host(), var_handle.relaxed(redux_op))->*[&](cudaStream_t s, auto var) {
+      cuda_safe_call(cudaStreamSynchronize(s));
+      *var.data_handle() += i;
+    };
+  }
+
+  // Check result
+  ctx.task(exec_place::host(), var_handle.read())->*[&](cudaStream_t s, auto var) {
+    cuda_safe_call(cudaStreamSynchronize(s));
+    int value    = *var.data_handle();
+    int expected = (N * (N - 1)) / 2 * (ndevs + 1);
+    EXPECT(value == expected);
+  };
+
+  ctx.finalize();
+}

cccl_upstream/cudax/test/stf/reductions/write_back_after_redux.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.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief This test ensures that data are properly reconstructed after a
+ *        reduction phase during the write-back mechanism
+ */
+
+#include <cuda/experimental/__stf/stream/interfaces/slice_reduction_ops.cuh>
+#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
+
+#include <iostream>
+
+using namespace cuda::experimental::stf;
+
+template <typename T>
+__global__ void add_val(slice<T> inout, T val)
+{
+  inout(0) += val;
+}
+
+int main()
+{
+  stream_ctx ctx;
+
+  int ndevs;
+  cuda_safe_call(cudaGetDeviceCount(&ndevs));
+
+  const int N = 4;
+
+  int var = 42;
+
+  auto var_handle = ctx.logical_data(make_slice(&var, 1));
+  var_handle.set_symbol("var");
+
+  auto redux_op = std::make_shared<slice_reduction_op_sum<int>>();
+
+  // We add i twice (total = N(N-1) + initial_value)
+  for (int i = 0; i < N; i++)
+  {
+    // device
+    for (int d = 0; d < ndevs; d++)
+    {
+      ctx.task(exec_place::device(d), var_handle.relaxed(redux_op))->*[=](cudaStream_t s, auto var) {
+        add_val<int><<<1, 1, 0, s>>>(var, i);
+      };
+    }
+
+    // host
+    ctx.host_launch(var_handle.relaxed(redux_op))->*[=](auto var) {
+      var(0) += i;
+    };
+  }
+
+  ctx.finalize();
+
+  int expected = 42 + (N * (N - 1)) / 2 * (ndevs + 1);
+  EXPECT(var == expected);
+}