[c.parallel]: migrate transform to use jit templates instead of string based implementations (#7399)

* Plumb custom storage_t through jit templates to handle cases where operators can have different storage types for input and output

* Set value_type of output iterator template to actual type instead of void to ensure proper policy selection in device code

* Disable well known operations for storage types

Author: NaderAlAwar

c/parallel/src/jit_templates/README.md

@@ -55,6 +55,9 @@ void func(/*..., */ cccl_iterator_t input_it /*, ...*/) {
             // `get_specialization` performs rudimentary type checking; if you pass the wrong type of an argument here
             // (one that doesn't match the types expected by the template), this call to `get_specialization` will fail to
             // compile.
+            //
+            // If different arguments require different storage types (e.g. input vs output CCCL_STORAGE), wrap them in
+            // tagged_arg<StorageT, T> so parameter_mapping can select the correct storage type per argument.
             input_t
         );
 
@@ -74,6 +77,10 @@ void func(/*..., */ cccl_iterator_t input_it /*, ...*/) {
 An argument mapping is a type, usable as a template argument, which carries the `cccl_*` values that are the
 arguments to the invocation of `get_specialization` in a format that will work in device code compiled with NVRTC.
 
+If different runtime arguments require different storage types, wrap the runtime argument in
+`tagged_arg<StorageT, T>` (defined in `traits.h`). This allows `parameter_mapping` to select the correct storage type
+per argument.
+
 Below is the contents of the `cccl_type_info` mapping header with annotations:
 
 ```cpp
@@ -162,6 +169,8 @@ A template traits type must provide:
    given arguments. The poster child for the use of this feature is the iterator JIT templates, where, if the kind of
    the iterator is determined to be a pointer, a simple pointer name is returned (instead of generating a specialization
    name of the template).
+   Note: if callers wrap runtime arguments in `tagged_arg`, traits that want to participate in `special` handling
+   should provide matching overloads.
 
 #### Archetypes
 
@@ -180,6 +189,9 @@ information provided by template traits and archetypes; performs basic type chec
 plus any auxiliary code necessary to compile it. See the example at the beginning of this document for an explanation of
 its various parameters and the return values.
 
+`get_specialization` will call `Traits::special(...)` when that call is well-formed for the provided runtime arguments,
+including tagged arguments.
+
 ### CMake
 
 The last bit necessary to make all of this work is the preprocessing and embedding step. In it, the `jit_entry.h` header

c/parallel/src/jit_templates/mappings/iterator.h

@@ -37,52 +37,61 @@ struct parameter_mapping<cccl_iterator_t>
 {
   static const constexpr auto archetype = cccl_iterator_t_mapping<int>{};
 
-  template <typename Traits>
-  static std::string map(template_id<Traits>, cccl_iterator_t arg)
+  template <typename Traits, typename ArgT>
+  static std::string map(template_id<Traits>, ArgT arg)
   {
-    if (arg.advance.type != cccl_op_kind_t::CCCL_STATEFUL && arg.advance.type != cccl_op_kind_t::CCCL_STATELESS)
+    using traits       = arg_traits<cuda::std::decay_t<ArgT>>;
+    using storage_type = typename traits::storage_type;
+    const auto& value  = traits::unwrap(arg);
+
+    if (value.advance.type != cccl_op_kind_t::CCCL_STATEFUL && value.advance.type != cccl_op_kind_t::CCCL_STATELESS)
     {
       throw std::runtime_error("c.parallel: well-known operations are not allowed as an iterator's advance operation");
     }
-    if (arg.dereference.type != cccl_op_kind_t::CCCL_STATEFUL && arg.dereference.type != cccl_op_kind_t::CCCL_STATELESS)
+    if (value.dereference.type != cccl_op_kind_t::CCCL_STATEFUL
+        && value.dereference.type != cccl_op_kind_t::CCCL_STATELESS)
     {
       throw std::runtime_error("c.parallel: well-known operations are not allowed as an iterator's dereference "
                                "operation");
     }
 
     return std::format(
       "cccl_iterator_t_mapping<{}>{{.is_pointer = {}, .size = {}, .alignment = {}, .advance = {}, .{} = {}}}",
-      cccl_type_enum_to_name(arg.value_type.type),
-      arg.type == cccl_iterator_kind_t::CCCL_POINTER,
-      arg.size,
-      arg.alignment,
-      arg.advance.name,
+      cccl_type_enum_to_name<storage_type>(value.value_type.type),
+      value.type == cccl_iterator_kind_t::CCCL_POINTER,
+      value.size,
+      value.alignment,
+      value.advance.name,
       std::is_same_v<Traits, output_iterator_traits> ? "assign" : "dereference",
-      arg.dereference.name);
+      value.dereference.name);
   }
 
-  template <typename Traits>
-  static std::string aux(template_id<Traits>, cccl_iterator_t arg)
+  template <typename Traits, typename ArgT>
+  static std::string aux(template_id<Traits>, ArgT arg)
   {
+    using traits       = arg_traits<cuda::std::decay_t<ArgT>>;
+    using storage_type = typename traits::storage_type;
+    const auto& value  = traits::unwrap(arg);
+
     if constexpr (std::is_same_v<Traits, output_iterator_traits>)
     {
       return std::format(
         R"output(
 extern "C" __device__ void {0}(void *, const void*);
 extern "C" __device__ void {1}(const void *, const void*);
 )output",
-        arg.advance.name,
-        arg.dereference.name);
+        value.advance.name,
+        value.dereference.name);
     }
 
     return std::format(
       R"input(
 extern "C" __device__ void {0}(void *, const void*);
 extern "C" __device__ void {1}(const void *, {2}*);
 )input",
-      arg.advance.name,
-      arg.dereference.name,
-      cccl_type_enum_to_name(arg.value_type.type));
+      value.advance.name,
+      value.dereference.name,
+      cccl_type_enum_to_name<storage_type>(value.value_type.type));
   }
 };
 #endif

c/parallel/src/jit_templates/mappings/operation.h

@@ -32,24 +32,26 @@ struct parameter_mapping<cccl_op_t>
 {
   static const constexpr auto archetype = cccl_op_t_mapping{};
 
-  template <typename Traits>
-  static std::string map(template_id<Traits>, cccl_op_t op)
+  template <typename Traits, typename ArgT>
+  static std::string map(template_id<Traits>, ArgT arg)
   {
+    const auto& value = arg_traits<cuda::std::decay_t<ArgT>>::unwrap(arg);
     return std::format(
       "cccl_op_t_mapping{{.is_stateless = {}, .size = {}, .alignment = {}, .operation = {}}}",
-      op.type != cccl_op_kind_t::CCCL_STATEFUL,
-      op.size,
-      op.alignment,
-      op.name);
+      value.type != cccl_op_kind_t::CCCL_STATEFUL,
+      value.size,
+      value.alignment,
+      value.name);
   }
 
-  template <typename Traits>
-  static std::string aux(template_id<Traits>, cccl_op_t op)
+  template <typename Traits, typename ArgT>
+  static std::string aux(template_id<Traits>, ArgT arg)
   {
+    const auto& value = arg_traits<cuda::std::decay_t<ArgT>>::unwrap(arg);
     return std::format(R"(
         extern "C" __device__ void {}();
         )",
-                       op.name);
+                       value.name);
   }
 };
 #endif

c/parallel/src/jit_templates/mappings/type_info.h

@@ -24,14 +24,17 @@ struct parameter_mapping<cccl_type_info>
 {
   static const constexpr auto archetype = cccl_type_info_mapping<int>{};
 
-  template <typename TplId>
-  static std::string map(TplId, cccl_type_info arg)
+  template <typename TplId, typename ArgT>
+  static std::string map(TplId, ArgT arg)
   {
-    return std::format("cccl_type_info_mapping<{}>{{}}", cccl_type_enum_to_name(arg.type));
+    using traits       = arg_traits<cuda::std::decay_t<ArgT>>;
+    using storage_type = typename traits::storage_type;
+    const auto& value  = traits::unwrap(arg);
+    return std::format("cccl_type_info_mapping<{}>{{}}", cccl_type_enum_to_name<storage_type>(value.type));
   }
 
-  template <typename TplId>
-  static std::string aux(TplId, cccl_type_info)
+  template <typename TplId, typename ArgT>
+  static std::string aux(TplId, ArgT)
   {
     return {};
   }

c/parallel/src/jit_templates/templates/input_iterator.h

@@ -84,5 +84,17 @@ struct input_iterator_traits
 
     return cuda::std::nullopt;
   }
+
+  template <typename Tag, typename StorageT>
+  static cuda::std::optional<specialization> special(tagged_arg<StorageT, cccl_iterator_t> it)
+  {
+    if (it.value.type == cccl_iterator_kind_t::CCCL_POINTER)
+    {
+      return cuda::std::make_optional(
+        specialization{cccl_type_enum_to_name<StorageT>(it.value.value_type.type, true), ""});
+    }
+
+    return cuda::std::nullopt;
+  }
 };
 #endif

c/parallel/src/jit_templates/templates/operation.h

@@ -289,36 +289,67 @@ __device__ {1} operator{2}(const {3} & lhs, const {3} & rhs)
     }
   }
 
-  template <typename, typename... Args>
-  static cuda::std::optional<specialization> special(cccl_op_t operation, cccl_type_info ret, Args... arguments)
+  template <typename Tag, typename RetStorageT, typename... ArgStorageTs>
+  static cuda::std::optional<specialization>
+  special(cccl_op_t operation,
+          tagged_arg<RetStorageT, cccl_type_info> ret,
+          tagged_arg<ArgStorageTs, cccl_type_info>... arguments)
   {
+    // We cannot use well-known operations with storage types as there
+    // is currently no way to tell whether multiple storage types,
+    // e.g., input and output, are the same type. This is necessary in
+    // order for operations like `negate` to work, as it requires both
+    // input and output to be the same type. For now, this check short
+    // circuits the specialization process for well-known operations
+    // with storage types. The code below that checks whether any of
+    // the arguments or the return type are storage types will
+    // currently not run, but is left here as it will be needed in the
+    // future.
+    if (ret.value.type == cccl_type_enum::CCCL_STORAGE
+        || ((arguments.value.type == cccl_type_enum::CCCL_STORAGE) || ...))
+    {
+      return cuda::std::nullopt;
+    }
+
     auto entry = well_known_operation_description(operation.type);
     if (!entry)
     {
       return cuda::std::nullopt;
     }
 
-    cccl_type_enum type_info_table[] = {ret.type, arguments.type...};
+    cccl_type_enum type_info_table[] = {ret.value.type, arguments.value.type...};
     auto builder                     = entry->check(cuda::std::span(type_info_table), std::format(entry->name, +""));
 
     std::string aux = "#include <cuda/std/functional>\n";
     if (entry->symbol)
     {
-      if (((arguments.type == cccl_type_enum::CCCL_STORAGE) || ...))
+      if (((arguments.value.type == cccl_type_enum::CCCL_STORAGE) || ...))
       {
-        std::string type_names[] = {cccl_type_enum_to_name(ret.type), cccl_type_enum_to_name(arguments.type)...};
+        std::string type_names[] = {cccl_type_enum_to_name<RetStorageT>(ret.value.type),
+                                    cccl_type_enum_to_name<ArgStorageTs>(arguments.value.type)...};
         auto type_name_views     = [&]<auto... Is>(std::index_sequence<Is...>) {
           return std::array<std::string_view, sizeof...(Is)>{{type_names[Is]...}};
         }(std::make_index_sequence<1 + sizeof...(arguments)>());
         aux += builder(cuda::std::span(type_name_views), *entry->symbol, operation.name);
       }
-      else if (ret.type == cccl_type_enum::CCCL_STORAGE)
+      else if (ret.value.type == cccl_type_enum::CCCL_STORAGE)
       {
         return cuda::std::nullopt;
       }
     }
 
-    return specialization{std::format(entry->name, cccl_type_enum_to_name(type_info_table[1]).c_str()), aux};
+    // Format the specialization name using the appropriate storage type
+    // type_info_table[1] corresponds to the first argument
+    using FirstArgStorageT = typename cuda::std::tuple_element<0, cuda::std::tuple<ArgStorageTs...>>::type;
+    std::string type_name  = cccl_type_enum_to_name<FirstArgStorageT>(type_info_table[1]);
+
+    return specialization{std::format(entry->name, type_name.c_str()), aux};
+  }
+
+  template <typename Tag, typename... Args>
+  static cuda::std::optional<specialization> special(cccl_op_t operation, Args... args)
+  {
+    return special<Tag>(operation, arg_traits<cuda::std::decay_t<Args>>::wrap(args)...);
   }
 #endif
 };
@@ -330,10 +361,11 @@ struct binary_user_operation_traits
   using type = user_operation_traits::type<Tag, Operation, ValueT, ValueT, ValueT>;
 
 #ifndef _CCCL_C_PARALLEL_JIT_TEMPLATES_PREPROCESS
-  template <typename Tag, typename... Args>
-  static cuda::std::optional<specialization> special(cccl_op_t operation, cccl_type_info arg_t)
+  template <typename Tag, typename Arg>
+  static cuda::std::optional<specialization> special(cccl_op_t operation, Arg arg)
   {
-    return user_operation_traits::special<Tag>(operation, arg_t, arg_t, arg_t);
+    auto wrapped = arg_traits<cuda::std::decay_t<Arg>>::wrap(arg);
+    return user_operation_traits::special<Tag>(operation, wrapped, wrapped, wrapped);
   }
 #endif
 };
@@ -345,10 +377,11 @@ struct binary_user_predicate_traits
   using type = user_operation_traits::type<Tag, Operation, cccl_type_info_mapping<bool>{}, ValueT, ValueT>;
 
 #ifndef _CCCL_C_PARALLEL_JIT_TEMPLATES_PREPROCESS
-  template <typename Tag, typename... Args>
-  static cuda::std::optional<specialization> special(cccl_op_t operation, cccl_type_info arg_t)
+  template <typename Tag, typename Arg>
+  static cuda::std::optional<specialization> special(cccl_op_t operation, Arg arg)
   {
-    return user_operation_traits::special<Tag>(operation, arg_t, arg_t, arg_t);
+    auto wrapped = arg_traits<cuda::std::decay_t<Arg>>::wrap(arg);
+    return user_operation_traits::special<Tag>(operation, wrapped, wrapped, wrapped);
   }
 #endif
 };

c/parallel/src/jit_templates/templates/output_iterator.h

@@ -47,7 +47,7 @@ struct output_iterator_t
 {
   using iterator_category = cuda::std::random_access_iterator_tag;
   using difference_type   = cuda::std::size_t;
-  using value_type        = void;
+  using value_type        = typename decltype(AssignTV)::Type;
   using reference =
     output_iterator_proxy_t<Tag, Iterator.size, Iterator.alignment, typename decltype(AssignTV)::Type, Iterator.assign>;
   using pointer = reference*;
@@ -98,5 +98,18 @@ struct output_iterator_traits
 
     return cuda::std::nullopt;
   }
+
+  template <typename Tag, typename StorageT, typename AssignStorageT>
+  static cuda::std::optional<specialization>
+  special(tagged_arg<StorageT, cccl_iterator_t> it, tagged_arg<AssignStorageT, cccl_type_info> assign_t)
+  {
+    if (it.value.type == cccl_iterator_kind_t::CCCL_POINTER)
+    {
+      return cuda::std::make_optional(
+        specialization{cccl_type_enum_to_name<AssignStorageT>(assign_t.value.type, true), ""});
+    }
+
+    return cuda::std::nullopt;
+  }
 };
 #endif