[STF] Support compound while-loop conditions in the C API

c/experimental/stf/include/cccl/c/experimental/stf/stf.h

@@ -2050,6 +2050,51 @@ void stf_stackable_while_cond_scalar(
   double threshold,
   stf_dtype dtype);
 
+//! \brief Maximum number of terms accepted by \c stf_stackable_while_cond_multi().
+#  define STF_WHILE_COND_MAX_TERMS 8
+
+//! \brief Combiner for multi-term while conditions.
+typedef enum stf_cond_combiner
+{
+  STF_COND_ALL = 0, //!< Continue while every term holds (logical AND)
+  STF_COND_ANY = 1, //!< Continue while at least one term holds (logical OR)
+} stf_cond_combiner;
+
+//! \brief One comparison term of a multi-term while condition.
+//!
+//! Evaluates as ``(*ld <op> threshold)``, then negated if \c negate is nonzero.
+typedef struct stf_while_cond_term
+{
+  stf_logical_data_handle ld; //!< Scalar logical data (1 element of \c dtype)
+  stf_compare_op op; //!< Comparison operator
+  double threshold; //!< Right-hand side compared against the scalar
+  stf_dtype dtype; //!< Element type of \c ld
+  int negate; //!< Nonzero to negate the term result
+} stf_while_cond_term;
+
+//! \brief Set a compound while-loop condition combining several scalar comparisons.
+//!
+//! Schedules a single internal task that reads every referenced scalar
+//! logical data, evaluates each ``(*ld <op> threshold)`` term (with optional
+//! negation) and folds the results with the requested combiner: continue
+//! while *all* terms hold (\c STF_COND_ALL) or while *any* term holds
+//! (\c STF_COND_ANY). Call exactly once per iteration after the loop body
+//! tasks of the current scope.
+//!
+//! With \p n_terms == 1 this is equivalent to \c stf_stackable_while_cond_scalar().
+//!
+//! \param ctx      Stackable context handle
+//! \param scope    While scope handle
+//! \param terms    Array of \p n_terms comparison terms
+//! \param n_terms  Number of terms (1 to \c STF_WHILE_COND_MAX_TERMS)
+//! \param combiner How the term results are combined
+void stf_stackable_while_cond_multi(
+  stf_ctx_handle ctx,
+  stf_while_scope_handle scope,
+  const stf_while_cond_term* terms,
+  int n_terms,
+  stf_cond_combiner combiner);
+
 #endif // CUDART_VERSION >= 12040
 
 //! \brief Create stackable logical data from existing memory and a data place.

c/experimental/stf/src/stf.cu

@@ -1279,33 +1279,77 @@ decltype(auto) visit_sld(stf_logical_data_handle h, F&& f)
 }
 
 #if _CCCL_CTK_AT_LEAST(12, 4)
-// Built-in condition kernel for while_cond_scalar.  Reads the head of the
-// scalar logical data, applies the requested comparison and updates the
-// conditional handle in place.  Lives outside extern "C" because it is a
-// device kernel template.
-template <typename T>
-__global__ void
-stf_stackable_while_cond_kernel(const T* value, cudaGraphConditionalHandle handle, double threshold, int op)
-{
-  const double v = static_cast<double>(*value);
-  bool result;
-  switch (op)
+// Built-in condition kernel for while_cond_scalar / while_cond_multi.  Reads
+// the head scalar of each referenced logical data, applies the requested
+// comparison (optionally negated), folds the term results with the requested
+// combiner and updates the conditional handle in place.  The whole term pack
+// is passed by value through kernel parameters (at most
+// STF_WHILE_COND_MAX_TERMS entries, well under the parameter-space limit) so
+// no device allocation is needed.  Lives outside extern "C" because it is a
+// device kernel.
+struct stf_while_cond_term_dev
+{
+  const void* ptr;
+  double threshold;
+  int op;
+  int dtype;
+  int negate;
+};
+
+struct stf_while_cond_pack
+{
+  stf_while_cond_term_dev terms[STF_WHILE_COND_MAX_TERMS];
+  int n_terms;
+  int combiner;
+};
+
+__global__ void stf_stackable_while_cond_kernel(stf_while_cond_pack pack, cudaGraphConditionalHandle handle)
+{
+  bool result = (pack.combiner == STF_COND_ALL);
+  for (int i = 0; i < pack.n_terms; ++i)
   {
-    case STF_CMP_GT:
-      result = v > threshold;
-      break;
-    case STF_CMP_LT:
-      result = v < threshold;
-      break;
-    case STF_CMP_GE:
-      result = v >= threshold;
-      break;
-    case STF_CMP_LE:
-      result = v <= threshold;
-      break;
-    default:
-      result = false;
-      break;
+    const stf_while_cond_term_dev& t = pack.terms[i];
+    double v                         = 0.0;
+    switch (t.dtype)
+    {
+      case STF_DTYPE_FLOAT32:
+        v = static_cast<double>(*static_cast<const float*>(t.ptr));
+        break;
+      case STF_DTYPE_FLOAT64:
+        v = *static_cast<const double*>(t.ptr);
+        break;
+      case STF_DTYPE_INT32:
+        v = static_cast<double>(*static_cast<const int*>(t.ptr));
+        break;
+      case STF_DTYPE_INT64:
+        v = static_cast<double>(*static_cast<const long long*>(t.ptr));
+        break;
+      default:
+        break;
+    }
+    bool term = false;
+    switch (t.op)
+    {
+      case STF_CMP_GT:
+        term = v > t.threshold;
+        break;
+      case STF_CMP_LT:
+        term = v < t.threshold;
+        break;
+      case STF_CMP_GE:
+        term = v >= t.threshold;
+        break;
+      case STF_CMP_LE:
+        term = v <= t.threshold;
+        break;
+      default:
+        break;
+    }
+    if (t.negate)
+    {
+      term = !term;
+    }
+    result = (pack.combiner == STF_COND_ALL) ? (result && term) : (result || term);
   }
   cudaGraphSetConditional(handle, result ? 1 : 0);
 }
@@ -1520,62 +1564,91 @@ void stf_stackable_while_cond_scalar(
   stf_compare_op op,
   double threshold,
   stf_dtype dtype)
+{
+  stf_while_cond_term term{ld, op, threshold, dtype, /* negate */ 0};
+  stf_stackable_while_cond_multi(ctx, scope, &term, 1, STF_COND_ALL);
+}
+
+void stf_stackable_while_cond_multi(
+  stf_ctx_handle ctx,
+  stf_while_scope_handle scope,
+  const stf_while_cond_term* terms,
+  int n_terms,
+  stf_cond_combiner combiner)
 {
   _CCCL_ASSERT(ctx != nullptr, "stackable context handle must not be null");
   _CCCL_ASSERT(scope != nullptr, "while scope handle must not be null");
-  _CCCL_ASSERT(ld != nullptr, "stackable logical data handle must not be null");
+  _CCCL_ASSERT(terms != nullptr, "condition terms must not be null");
+  _CCCL_ASSERT(n_terms >= 1 && n_terms <= STF_WHILE_COND_MAX_TERMS, "invalid number of condition terms");
+  _CCCL_ASSERT(combiner == STF_COND_ALL || combiner == STF_COND_ANY, "invalid condition combiner");
 
   auto* sctx                             = from_opaque_sctx(ctx);
   auto* guard                            = from_opaque_while(scope);
   cudaGraphConditionalHandle cond_handle = guard->cond_handle();
 
   const int offset = sctx->get_head_offset();
 
-  // Validate (and auto-push if necessary) the read access on this scope,
-  // then materialise the untyped logical_data for the task dep.  The
-  // concrete stackable_logical_data<T> is dispatched through visit_sld()
-  // so both slice<char>-backed data and void_interface tokens resolve
-  // correctly; the while-condition kernel below only makes sense on a
-  // scalar-typed slice, but validate_access/get_ld are type-agnostic.
-  logical_data_untyped ld_ut = visit_sld(ld, [&](auto& sld) {
-    sld.validate_access(offset, *sctx, access_mode::read);
-    return logical_data_untyped{sld.get_ld(offset)};
-  });
-
   auto& underlying_ctx = sctx->get_ctx(offset);
   auto task            = underlying_ctx.task();
-  task.add_deps(task_dep_untyped(ld_ut, access_mode::read));
+
+  // Validate (and auto-push if necessary) the read access of every term on
+  // this scope, then materialise the untyped logical_data for the task dep.
+  // The concrete stackable_logical_data<T> is dispatched through visit_sld()
+  // so both slice<char>-backed data and void_interface tokens resolve
+  // correctly; the while-condition kernel below only makes sense on
+  // scalar-typed slices, but validate_access/get_ld are type-agnostic.
+  // Duplicate handles share a single read dependency (dep_index remembers
+  // which task-dep slot each term resolves to).
+  int dep_index[STF_WHILE_COND_MAX_TERMS];
+  int n_deps = 0;
+  for (int i = 0; i < n_terms; ++i)
+  {
+    _CCCL_ASSERT(terms[i].ld != nullptr, "stackable logical data handle must not be null");
+    _CCCL_ASSERT(terms[i].dtype >= STF_DTYPE_FLOAT32 && terms[i].dtype <= STF_DTYPE_INT64,
+                 "unsupported dtype for stf_stackable_while_cond_multi");
+    int found = -1;
+    for (int j = 0; j < i; ++j)
+    {
+      if (terms[j].ld == terms[i].ld)
+      {
+        found = dep_index[j];
+        break;
+      }
+    }
+    if (found >= 0)
+    {
+      dep_index[i] = found;
+      continue;
+    }
+    logical_data_untyped ld_ut = visit_sld(terms[i].ld, [&](auto& sld) {
+      sld.validate_access(offset, *sctx, access_mode::read);
+      return logical_data_untyped{sld.get_ld(offset)};
+    });
+    task.add_deps(task_dep_untyped(ld_ut, access_mode::read));
+    dep_index[i] = n_deps++;
+  }
+
   task.set_symbol("while_condition");
   task.enable_capture();
   task.start();
 
-  auto stream     = task.get_stream();
-  auto s          = task.template get<slice<const char>>(0);
-  const void* ptr = s.data_handle();
+  auto stream = task.get_stream();
 
-  switch (dtype)
+  stf_while_cond_pack pack{};
+  pack.n_terms  = n_terms;
+  pack.combiner = static_cast<int>(combiner);
+  for (int i = 0; i < n_terms; ++i)
   {
-    case STF_DTYPE_FLOAT32:
-      stf_stackable_while_cond_kernel<float>
-        <<<1, 1, 0, stream>>>(static_cast<const float*>(ptr), cond_handle, threshold, op);
-      break;
-    case STF_DTYPE_FLOAT64:
-      stf_stackable_while_cond_kernel<double>
-        <<<1, 1, 0, stream>>>(static_cast<const double*>(ptr), cond_handle, threshold, op);
-      break;
-    case STF_DTYPE_INT32:
-      stf_stackable_while_cond_kernel<int>
-        <<<1, 1, 0, stream>>>(static_cast<const int*>(ptr), cond_handle, threshold, op);
-      break;
-    case STF_DTYPE_INT64:
-      stf_stackable_while_cond_kernel<long long>
-        <<<1, 1, 0, stream>>>(static_cast<const long long*>(ptr), cond_handle, threshold, op);
-      break;
-    default:
-      _CCCL_ASSERT(false, "unsupported dtype for stf_stackable_while_cond_scalar");
-      break;
+    auto s                  = task.template get<slice<const char>>(dep_index[i]);
+    pack.terms[i].ptr       = s.data_handle();
+    pack.terms[i].threshold = terms[i].threshold;
+    pack.terms[i].op        = static_cast<int>(terms[i].op);
+    pack.terms[i].dtype     = static_cast<int>(terms[i].dtype);
+    pack.terms[i].negate    = terms[i].negate;
   }
 
+  stf_stackable_while_cond_kernel<<<1, 1, 0, stream>>>(pack, cond_handle);
+
   task.end();
 }
 

c/experimental/stf/test/test_stackable.cu

@@ -659,4 +659,103 @@ C2H_TEST("stackable: while-body K chained rw tasks sweep", "[stackable][while][c
   (void) total_off_by_one;
 }
 
+namespace
+{
+// Run a while loop whose body increments a 1-element double counter once per
+// iteration and leaves a 1-element flag at its initial value 1.0.  The
+// continuation condition is built by `set_condition` from the counter and
+// flag handles.  Returns the final counter value observed on the host.
+template <typename SetCondition>
+double run_compound_while(SetCondition&& set_condition)
+{
+  stf_ctx_handle ctx = stf_stackable_ctx_create();
+  REQUIRE(ctx != nullptr);
+
+  double* host_iter;
+  REQUIRE(cudaMallocHost(&host_iter, sizeof(double)) == cudaSuccess);
+  host_iter[0] = 0.0;
+  double* host_flag;
+  REQUIRE(cudaMallocHost(&host_flag, sizeof(double)) == cudaSuccess);
+  host_flag[0] = 1.0;
+
+  stf_logical_data_handle lIter = stf_stackable_logical_data(ctx, host_iter, sizeof(double));
+  REQUIRE(lIter != nullptr);
+  stf_logical_data_handle lFlag = stf_stackable_logical_data(ctx, host_flag, sizeof(double));
+  REQUIRE(lFlag != nullptr);
+
+  stf_while_scope_handle scope = stf_stackable_push_while(ctx);
+  REQUIRE(scope != nullptr);
+  {
+    stf_task_handle t = stf_stackable_task_create(ctx);
+    REQUIRE(t != nullptr);
+    stf_stackable_task_add_dep(ctx, t, lIter, STF_RW);
+    stf_task_enable_capture(t);
+    stf_task_start(t);
+    double* d = static_cast<double*>(stf_task_get(t, 0));
+    increment_kernel<<<1, 1, 0, (cudaStream_t) stf_task_get_custream(t)>>>(1, d);
+    stf_task_end(t);
+    stf_task_destroy(t);
+
+    set_condition(ctx, scope, lIter, lFlag);
+  }
+  stf_stackable_pop_while(scope);
+
+  stf_stackable_logical_data_destroy(lIter);
+  stf_stackable_logical_data_destroy(lFlag);
+  stf_stackable_ctx_finalize(ctx);
+
+  const double result = host_iter[0];
+  REQUIRE(cudaFreeHost(host_iter) == cudaSuccess);
+  REQUIRE(cudaFreeHost(host_flag) == cudaSuccess);
+  return result;
+}
+} // namespace
+
+C2H_TEST("stackable: while compound condition", "[stackable][while][c-api]")
+{
+  SECTION("ALL combiner stops at the iteration cap")
+  {
+    // flag > 0.5 is always true; iter < 5 caps the loop at 5 iterations.
+    const double iters = run_compound_while(
+      [](stf_ctx_handle ctx, stf_while_scope_handle scope, stf_logical_data_handle lIter, stf_logical_data_handle lFlag) {
+        stf_while_cond_term terms[2] = {
+          {lFlag, STF_CMP_GT, 0.5, STF_DTYPE_FLOAT64, 0},
+          {lIter, STF_CMP_LT, 5.0, STF_DTYPE_FLOAT64, 0},
+        };
+        stf_stackable_while_cond_multi(ctx, scope, terms, 2, STF_COND_ALL);
+      });
+    REQUIRE(iters == 5.0);
+  }
+
+  SECTION("ANY combiner with a negated term")
+  {
+    // ~(flag > 0.5) is always false, so only iter < 3 keeps the loop going.
+    const double iters = run_compound_while(
+      [](stf_ctx_handle ctx, stf_while_scope_handle scope, stf_logical_data_handle lIter, stf_logical_data_handle lFlag) {
+        stf_while_cond_term terms[2] = {
+          {lIter, STF_CMP_LT, 3.0, STF_DTYPE_FLOAT64, 0},
+          {lFlag, STF_CMP_GT, 0.5, STF_DTYPE_FLOAT64, 1},
+        };
+        stf_stackable_while_cond_multi(ctx, scope, terms, 2, STF_COND_ANY);
+      });
+    REQUIRE(iters == 3.0);
+  }
+
+  SECTION("duplicate logical data across terms shares one dependency")
+  {
+    const double iters = run_compound_while(
+      [](stf_ctx_handle ctx,
+         stf_while_scope_handle scope,
+         stf_logical_data_handle lIter,
+         stf_logical_data_handle /*lFlag*/) {
+        stf_while_cond_term terms[2] = {
+          {lIter, STF_CMP_LT, 4.0, STF_DTYPE_FLOAT64, 0},
+          {lIter, STF_CMP_GT, -1.0, STF_DTYPE_FLOAT64, 0},
+        };
+        stf_stackable_while_cond_multi(ctx, scope, terms, 2, STF_COND_ALL);
+      });
+    REQUIRE(iters == 4.0);
+  }
+}
+
 #endif // _CCCL_CTK_AT_LEAST(12, 4)