Merge branch 'main' into stf_c_api

Author: caugonnet

cudax/examples/stf/CMakeLists.txt

@@ -21,28 +21,42 @@ set(
   01-axpy-launch.cu
   01-axpy-parallel_for.cu
   binary_fhe.cu
+  binary_fhe_stackable.cu
   09-dot-reduce.cu
   cfd.cu
   custom_data_interface.cu
   fdtd_mgpu.cu
+  fdtd_while.cu
+  fdtd_repeat_n.cu
   frozen_data_init.cu
   graph_algorithms/degree_centrality.cu
   graph_algorithms/jaccard.cu
   graph_algorithms/pagerank.cu
+  graph_algorithms/pagerank_batched.cu
+  graph_algorithms/pagerank_while.cu
   graph_algorithms/tricount.cu
+  graph_scope.cu
   heat.cu
   heat_mgpu.cu
   jacobi.cu
   jacobi_pfor.cu
+  jacobi_stackable.cu
+  jacobi_stackable_raii.cu
+  jacobi_update_cond.cu
   launch_histogram.cu
   launch_scan.cu
   launch_sum.cu
   launch_sum_cub.cu
+  linear_algebra/burger.cu
+  linear_algebra/burger_sensitivity.cu
+  linear_algebra/cg_csr.cu
+  linear_algebra/cg_csr_stackable.cu
   logical_gates_composition.cu
   mandelbrot.cu
   parallel_for_2D.cu
   pi.cu
   scan.cu
+  sqrt_newton_stackable.cu
   standalone-launches.cu
   word_count.cu
   word_count_reduce.cu
@@ -52,9 +66,9 @@ set(
 set(
   stf_example_mathlib_sources
   linear_algebra/06-pdgemm.cu
+  linear_algebra/06-pdgemm-stackable.cu
   linear_algebra/07-cholesky.cu
   linear_algebra/07-potri.cu
-  linear_algebra/cg_csr.cu
   linear_algebra/cg_dense_2D.cu
   linear_algebra/strassen.cu
 )

cudax/examples/stf/binary_fhe_stackable.cu

@@ -0,0 +1,240 @@
+//===----------------------------------------------------------------------===//
+//
+// 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) 2024-2025 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+/**
+ * @file
+ * @brief A toy example to illustrate how we can compose logical operations over encrypted data
+ */
+
+#include <cuda/experimental/stf.cuh>
+
+using namespace cuda::experimental::stf;
+
+#include <memory>
+
+class ciphertext;
+
+class plaintext
+{
+public:
+  plaintext(const stackable_ctx& ctx)
+      : ctx(ctx)
+  {}
+
+  plaintext(stackable_ctx& ctx, ::std::vector<char> v)
+      : values(mv(v))
+      , ctx(ctx)
+      , ld(ctx.logical_data(values.data(), values.size()))
+  {}
+
+  auto& set_symbol(const std::string& s)
+  {
+    ld.set_symbol(s);
+    symbol = s;
+
+    return *this;
+  }
+
+  const std::string& get_symbol() const
+  {
+    return symbol;
+  }
+
+  // This will asynchronously fill string s
+  void convert_to_vector(std::vector<char>& v)
+  {
+    ctx.host_launch(ld.read()).set_symbol("to_vector")->*[&](auto dl) {
+      v.resize(dl.size());
+      for (size_t i = 0; i < dl.size(); i++)
+      {
+        v[i] = dl(i);
+      }
+    };
+  }
+
+  ciphertext encrypt() const;
+
+private:
+  std::vector<char> values;
+  mutable stackable_ctx ctx;
+  ::std::string symbol;
+
+public:
+  mutable stackable_logical_data<slice<char>> ld;
+};
+
+class ciphertext
+{
+public:
+  ciphertext() = default;
+
+  // We need a deep-copy semantic
+  ciphertext(const ciphertext& other)
+      : ctx(other.ctx)
+      , symbol(other.symbol)
+  {
+    copy_content(ctx, other, *this);
+  }
+
+  ciphertext(const stackable_ctx& ctx)
+      : ctx(ctx)
+  {}
+
+  ciphertext(ciphertext&&)            = default;
+  ciphertext& operator=(ciphertext&&) = default;
+
+  static void copy_content(stackable_ctx& ctx, const ciphertext& src, ciphertext& dst)
+  {
+    dst.ld = ctx.logical_data(src.ld.shape());
+    ctx.parallel_for(src.ld.shape(), src.ld.read(), dst.ld.write()).set_symbol("copy")->*
+      [] __device__(size_t i, auto src, auto dst) {
+        dst(i) = src(i);
+      };
+  }
+
+  auto& set_symbol(std::string s)
+  {
+    ld.set_symbol(s);
+    symbol = mv(s);
+
+    return *this;
+  }
+
+  const std::string& get_symbol() const
+  {
+    return symbol;
+  }
+
+  plaintext decrypt() const
+  {
+    plaintext p(ctx);
+    p.ld = ctx.logical_data(shape_of<slice<char>>(ld.shape().size()));
+    ctx.parallel_for(ld.shape(), ld.read(), p.ld.write()).set_symbol("decrypt")->*
+      [] __device__(size_t i, auto cipher_data, auto plain_data) {
+        plain_data(i) = static_cast<char>(cipher_data(i) >> 32);
+      };
+    return p;
+  }
+
+  // Copy assignment operator
+  // We need a deep-copy semantic
+  ciphertext& operator=(const ciphertext& other)
+  {
+    if (this != &other)
+    {
+      ctx    = other.ctx;
+      symbol = other.symbol;
+      copy_content(ctx, other, *this);
+    }
+    return *this;
+  }
+
+  ciphertext operator|(const ciphertext& other) const
+  {
+    ciphertext result(ctx);
+    result.ld = ctx.logical_data(ld.shape());
+
+    ctx.parallel_for(ld.shape(), ld.read(), other.ld.read(), result.ld.write()).set_symbol("OR")->*
+      [] __device__(size_t i, auto d_c1, auto d_c2, auto d_res) {
+        d_res(i) = d_c1(i) | d_c2(i);
+      };
+
+    return result;
+  }
+
+  ciphertext operator&(const ciphertext& other) const
+  {
+    ciphertext result(ctx);
+    result.ld = ctx.logical_data(ld.shape());
+
+    ctx.parallel_for(ld.shape(), ld.read(), other.ld.read(), result.ld.write()).set_symbol("AND")->*
+      [] __device__(size_t i, auto d_c1, auto d_c2, auto d_res) {
+        d_res(i) = d_c1(i) & d_c2(i);
+      };
+
+    return result;
+  }
+
+  ciphertext operator~() const
+  {
+    ciphertext result(ctx);
+    result.ld = ctx.logical_data(ld.shape());
+
+    ctx.parallel_for(ld.shape(), ld.read(), result.ld.write()).set_symbol("NOT")->*
+      [] __device__(size_t i, auto d_c, auto d_res) {
+        d_res(i) = ~d_c(i);
+      };
+
+    return result;
+  }
+
+  mutable stackable_logical_data<slice<uint64_t>> ld;
+
+private:
+  mutable stackable_ctx ctx;
+  ::std::string symbol;
+};
+
+ciphertext plaintext::encrypt() const
+{
+  ciphertext c(ctx);
+  c.ld = ctx.logical_data(shape_of<slice<uint64_t>>(ld.shape().size()));
+
+  ctx.parallel_for(ld.shape(), ld.read(), c.ld.write()).set_symbol("encrypt")->*
+    [] __device__(size_t i, auto dptxt, auto dctxt) {
+      // A super safe encryption !
+      dctxt(i) = ((uint64_t) (dptxt(i)) << 32 | 0x4);
+    };
+
+  return c;
+}
+
+template <typename T>
+T circuit(const T& a, const T& b)
+{
+  return ~((a | ~b) & (~a | b));
+}
+
+int main()
+{
+  stackable_ctx ctx;
+
+  const std::vector<char> vA{3, 3, 2, 2, 17};
+  plaintext pA(ctx, std::vector<char>(vA));
+  pA.set_symbol("A");
+
+  const std::vector<char> vB{1, 7, 7, 7, 49};
+  plaintext pB(ctx, std::vector<char>(vB));
+  pB.set_symbol("B");
+
+  auto s_encrypt = ctx.dot_section("encrypt");
+  auto eA        = pA.encrypt().set_symbol("A");
+  auto eB        = pB.encrypt().set_symbol("B");
+  s_encrypt.end();
+
+  ctx.push();
+
+  auto s_circuit = ctx.dot_section("circuit");
+  auto out       = circuit(eA, eB);
+  s_circuit.end();
+
+  ctx.pop();
+
+  std::vector<char> v_out;
+  out.decrypt().convert_to_vector(v_out);
+
+  ctx.finalize();
+
+  for (size_t i = 0; i < v_out.size(); i++)
+  {
+    char expected = circuit(vA[i], vB[i]);
+    EXPECT(expected == v_out[i]);
+  }
+}