Skip to content

Extend and refactor transform overloads in CUDA system #5238

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bernhardmgruber merged 6 commits into from
Jul 17, 2025
Merged

Extend and refactor transform overloads in CUDA system #5238

Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
dfd38f0
Extend and refactor transform overloads in CUDA system
bernhardmgruber Jul 14, 2025
e5b06d6
Reviewer feedback
bernhardmgruber Jul 14, 2025
9c3bcc3
Unwrap thrust raw references
bernhardmgruber Jul 17, 2025
b43fbb4
Remove HD macro from overloads /wo exec policy
bernhardmgruber Jul 17, 2025
7dadd60
Fix signature
bernhardmgruber Jul 17, 2025
7645263
Fix template parameter name
bernhardmgruber Jul 17, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
238 changes: 154 additions & 84 deletions thrust/thrust/system/cuda/detail/transform.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,14 +41,15 @@

# include <cub/device/device_transform.cuh>

# include <thrust/distance.h>
# include <thrust/iterator/zip_iterator.h>
# include <thrust/system/cuda/detail/dispatch.h>
# include <thrust/system/cuda/detail/parallel_for.h>
# include <thrust/system/cuda/detail/util.h>
# include <thrust/zip_function.h>

# include <cuda/__functional/address_stability.h>
# include <cuda/std/__algorithm/transform.h>
Comment thread
bernhardmgruber marked this conversation as resolved.
Resolved
# include <cuda/std/__iterator/distance.h>
# include <cuda/std/cstdint>

THRUST_NAMESPACE_BEGIN
Expand Down Expand Up @@ -95,7 +96,7 @@ struct unary_transform_f
output[idx] = op(raw_reference_cast(input[idx]));
}
}
}; // struct unary_transform_stencil_f
};

template <class InputIt, class OutputIt, class TransformOp, class Predicate>
struct unary_transform_f<InputIt, OutputIt, no_stencil_tag, TransformOp, Predicate>
Expand All @@ -121,7 +122,7 @@ struct unary_transform_f<InputIt, OutputIt, no_stencil_tag, TransformOp, Predica
output[idx] = op(raw_reference_cast(input[idx]));
}
}
}; // struct unary_transform_f
};

template <class InputIt1, class InputIt2, class OutputIt, class StencilIt, class TransformOp, class Predicate>
struct binary_transform_f
Expand Down Expand Up @@ -152,7 +153,7 @@ struct binary_transform_f
output[idx] = op(raw_reference_cast(input1[idx]), raw_reference_cast(input2[idx]));
}
}
}; // struct binary_transform_stencil_f
};

template <class InputIt1, class InputIt2, class OutputIt, class TransformOp, class Predicate>
struct binary_transform_f<InputIt1, InputIt2, OutputIt, no_stencil_tag, TransformOp, Predicate>
Expand Down Expand Up @@ -181,7 +182,7 @@ struct binary_transform_f<InputIt1, InputIt2, OutputIt, no_stencil_tag, Transfor
output[idx] = op(raw_reference_cast(input1[idx]), raw_reference_cast(input2[idx]));
}
}
}; // struct binary_transform_f
};

// EAN 2024-10-04: when force-inlined, gcc's optimizer will generate bad code
// for this function:
Expand All @@ -201,9 +202,7 @@ OutputIt _CCCL_HOST_DEVICE inline unary(
}

using unary_transform_t = unary_transform_f<InputIt, OutputIt, StencilIt, TransformOp, Predicate>;

cuda_cub::parallel_for(policy, unary_transform_t(items, result, stencil, transform_op, predicate), num_items);

cuda_cub::parallel_for(policy, unary_transform_t{items, result, stencil, transform_op, predicate}, num_items);
return result + num_items;
}

Expand Down Expand Up @@ -233,10 +232,8 @@ OutputIt _CCCL_HOST_DEVICE inline binary(
}

using binary_transform_t = binary_transform_f<InputIt1, InputIt2, OutputIt, StencilIt, TransformOp, Predicate>;

cuda_cub::parallel_for(
policy, binary_transform_t(items1, items2, result, stencil, transform_op, predicate), num_items);

policy, binary_transform_t{items1, items2, result, stencil, transform_op, predicate}, num_items);
return result + num_items;
}

Expand Down Expand Up @@ -275,41 +272,88 @@ OutputIt THRUST_FUNCTION cub_transform_many(
return result + num_items;
}

template <typename... Ts, std::size_t... Is>
THRUST_FUNCTION auto convert_to_std_tuple(tuple<Ts...> t, ::cuda::std::index_sequence<Is...>)
-> ::cuda::std::tuple<Ts...>
{
return ::cuda::std::tuple<Ts...>{get<Is>(t)...};
}

// unwrap zip_iterator and zip_function into their underlying iterators so cub::DeviceTransform can optimize them
template <class Derived, class Offset, class... InputIts, class OutputIt, class TransformOp>
OutputIt THRUST_FUNCTION cub_transform_many(
execution_policy<Derived>& policy,
::cuda::std::tuple<zip_iterator<tuple<InputIts...>>> firsts,
::cuda::std::tuple<zip_iterator<::cuda::std::tuple<InputIts...>>> firsts,
OutputIt result,
Offset num_items,
zip_function<TransformOp> transform_op)
{
return cub_transform_many(
policy,
convert_to_std_tuple(get<0>(firsts).get_iterator_tuple(), ::cuda::std::index_sequence_for<InputIts...>{}),
result,
num_items,
transform_op.underlying_function());
policy, get<0>(firsts).get_iterator_tuple(), result, num_items, transform_op.underlying_function());
}
} // namespace __transform

//-------------------------
// Thrust API entry points
//-------------------------
template <typename F>
struct raw_reference_cast_args
{
mutable F f; // mutable to support non-const F::operator()

template <typename... Ts>
THRUST_FUNCTION decltype(auto) operator()(Ts&&... args) const
{
return f(raw_reference_cast(::cuda::std::forward<Ts>(args))...);
}
};
} // namespace __transform

//-------------------------
// one input data stream
//-------------------------

template <typename Derived, typename InputIt, typename OutputIt, typename TransformOp>
THRUST_FUNCTION OutputIt
transform(execution_policy<Derived>& policy, InputIt first, InputIt last, OutputIt result, TransformOp transform_op)
{
THRUST_CDP_DISPATCH(
(return __transform::cub_transform_many(
policy, ::cuda::std::make_tuple(first), result, ::cuda::std::distance(first, last), transform_op);),
(return ::cuda::std::transform(
first, last, result, __transform::raw_reference_cast_args<TransformOp>{transform_op});));
}

template <typename Derived, typename InputIt, typename OutputIt, typename TransformOp>
THRUST_FUNCTION OutputIt transform_n(
execution_policy<Derived>& policy,
InputIt first,
::cuda::std::iter_difference_t<InputIt> num_items,
OutputIt result,
TransformOp transform_op)
{
THRUST_CDP_DISPATCH(
(return __transform::cub_transform_many(policy, ::cuda::std::make_tuple(first), result, num_items, transform_op);),
(return ::cuda::std::transform(
first, first + num_items, result, __transform::raw_reference_cast_args<TransformOp>{transform_op});));
}

template <typename Derived, typename InputIt, typename OutputIt, typename TransformOp, typename Predicate>
THRUST_FUNCTION OutputIt transform_if(
execution_policy<Derived>& policy,
InputIt first,
InputIt last,
OutputIt result,
TransformOp transform_op,
Predicate predicate)
{
return __transform::unary(
policy, first, result, ::cuda::std::distance(first, last), __transform::no_stencil_tag{}, transform_op, predicate);
}

template <typename Derived, typename InputIt, typename OutputIt, typename TransformOp, typename Predicate>
THRUST_FUNCTION OutputIt transform_if_n(
execution_policy<Derived>& policy,
InputIt first,
::cuda::std::iter_difference_t<InputIt> count,
OutputIt result,
TransformOp transform_op,
Predicate predicate)
{
return __transform::unary(policy, first, result, count, __transform::no_stencil_tag{}, transform_op, predicate);
}

// one input data stream + stencil

template <class Derived, class InputIt, class OutputIt, class StencilInputIt, class TransformOp, class Predicate>
OutputIt THRUST_FUNCTION transform_if(
THRUST_FUNCTION OutputIt transform_if(
execution_policy<Derived>& policy,
InputIt first,
InputIt last,
Expand All @@ -318,85 +362,111 @@ OutputIt THRUST_FUNCTION transform_if(
TransformOp transform_op,
Predicate predicate)
{
using size_type = thrust::detail::it_difference_t<InputIt>;
size_type num_items = static_cast<size_type>(::cuda::std::distance(first, last));
return __transform::unary(policy, first, result, num_items, stencil, transform_op, predicate);
} // func transform_if
return __transform::unary(policy, first, result, ::cuda::std::distance(first, last), stencil, transform_op, predicate);
}

template <class Derived, class InputIt, class OutputIt, class TransformOp, class Predicate>
OutputIt THRUST_FUNCTION transform_if(
template <typename Derived,
typename InputIt,
typename StencilInputIt,
typename OutputIt,
typename TransformOp,
typename Predicate>
THRUST_FUNCTION OutputIt transform_if_n(
execution_policy<Derived>& policy,
InputIt first,
InputIt last,
::cuda::std::iter_difference_t<InputIt> num_items,
StencilInputIt stencil,
OutputIt result,
TransformOp transform_op,
Predicate predicate)
{
return cuda_cub::transform_if(policy, first, last, __transform::no_stencil_tag(), result, transform_op, predicate);
} // func transform_if
return __transform::unary(policy, first, result, num_items, stencil, transform_op, predicate);
}

template <class Derived, class InputIt, class OutputIt, class TransformOp>
OutputIt THRUST_FUNCTION
transform(execution_policy<Derived>& policy, InputIt first, InputIt last, OutputIt result, TransformOp transform_op)
// two input data streams

template <typename Derived, typename InputIt1, typename InputIt2, typename OutputIt, typename BinaryTransformOp>
THRUST_FUNCTION OutputIt transform(
execution_policy<Derived>& policy,
InputIt1 first1,
InputIt1 last1,
InputIt2 first2,
OutputIt result,
BinaryTransformOp transform_op)
{
THRUST_CDP_DISPATCH(
(using size_type = thrust::detail::it_difference_t<InputIt>;
const auto num_items = static_cast<size_type>(::cuda::std::distance(first, last));
return __transform::cub_transform_many(policy, ::cuda::std::make_tuple(first), result, num_items, transform_op);),
(while (first != last) {
*result = transform_op(raw_reference_cast(*first));
++first;
++result;
} return result;));
} // func transform

//-------------------------
// two input data streams
//-------------------------
(return __transform::cub_transform_many(
policy,
::cuda::std::make_tuple(first1, first2),
result,
::cuda::std::distance(first1, last1),
transform_op);),
(return ::cuda::std::transform(
first1, last1, first2, result, __transform::raw_reference_cast_args<BinaryTransformOp>{transform_op});));
}

template <class Derived,
class InputIt1,
class InputIt2,
class StencilInputIt,
class OutputIt,
class TransformOp,
class Predicate>
OutputIt THRUST_FUNCTION transform_if(
template <typename Derived, typename InputIt1, typename InputIt2, typename OutputIt, typename BinaryTransformOp>
THRUST_FUNCTION OutputIt transform_n(
execution_policy<Derived>& policy,
InputIt1 first1,
::cuda::std::iter_difference_t<InputIt1> num_items,
InputIt2 first2,
OutputIt result,
BinaryTransformOp transform_op)
{
THRUST_CDP_DISPATCH(
(return __transform::cub_transform_many(
policy, ::cuda::std::make_tuple(first1, first2), result, num_items, transform_op);),
(return ::cuda::std::transform(first1,
first1 + num_items,
first2,
result,
__transform::raw_reference_cast_args<BinaryTransformOp>{transform_op});));
}

// two input data streams + stencil

template <typename Derived,
typename InputIt1,
typename InputIt2,
typename StencilInputIt,
typename OutputIt,
typename BinaryTransformOp,
typename Predicate>
THRUST_FUNCTION OutputIt transform_if(
execution_policy<Derived>& policy,
InputIt1 first1,
InputIt1 last1,
InputIt2 first2,
StencilInputIt stencil,
OutputIt result,
TransformOp transform_op,
BinaryTransformOp transform_op,
Predicate predicate)
{
using size_type = thrust::detail::it_difference_t<InputIt1>;
size_type num_items = static_cast<size_type>(::cuda::std::distance(first1, last1));
return __transform::binary(policy, first1, first2, result, num_items, stencil, transform_op, predicate);
} // func transform_if
return __transform::binary(
policy, first1, first2, result, ::cuda::std::distance(first1, last1), stencil, transform_op, predicate);
}

template <class Derived, class InputIt1, class InputIt2, class OutputIt, class TransformOp>
OutputIt THRUST_FUNCTION transform(
template <typename Derived,
typename InputIt1,
typename InputIt2,
typename StencilInputIt,
typename OutputIt,
typename BinaryTransformOp,
typename Predicate>
THRUST_FUNCTION OutputIt transform_if_n(
execution_policy<Derived>& policy,
InputIt1 first1,
InputIt1 last1,
::cuda::std::iter_difference_t<InputIt1> num_items,
InputIt2 first2,
StencilInputIt stencil,
OutputIt result,
TransformOp transform_op)
BinaryTransformOp transform_op,
Predicate predicate)
{
THRUST_CDP_DISPATCH(
(using size_type = thrust::detail::it_difference_t<InputIt1>;
const auto num_items = static_cast<size_type>(::cuda::std::distance(first1, last1));
return __transform::cub_transform_many(
policy, ::cuda::std::make_tuple(first1, first2), result, num_items, transform_op);),
(while (first1 != last1) {
*result = transform_op(raw_reference_cast(*first1), raw_reference_cast(*first2));
++first1;
++first2;
++result;
} return result;));
return __transform::binary(policy, first1, first2, result, num_items, stencil, transform_op, predicate);
}

} // namespace cuda_cub

THRUST_NAMESPACE_END
Expand Down