[FEA] Unify square_t functors across benchmarks and tests

[rkothari3]

Description

closes #9619

Replaces locally-defined square_t structs with thrust::square<T>
where semantically equivalent (value-returning T operator()(const T& x)),
and consolidates the in-place variant into a shared square_inplace_t
in nvbench_helper.cuh.

Value-returning square_t to thrust::square<T>:

• cub/benchmarks/bench/transform_reduce/sum.cu
• thrust/benchmarks/bench/transform_reduce/sum.cu
• cub/test/catch2_test_device_transform_reduce.cu

In-place square_t<T> to square_inplace_t<T> (added to nvbench_helper.cuh):

• thrust/benchmarks/bench/for_each/basic.cu
• thrust/benchmarks/bench/for_each_n/basic.cu
• libcudacxx/benchmarks/bench/for_each/basic.cu
• libcudacxx/benchmarks/bench/for_each_n/basic.cu

Intentionally not changed:

• catch2_test_device_for_api.cu / for_env_api.cu: square_t is inside // example-begin documentation tags
• catch2_test_device_reduce_deterministic.cu: takes int input, different signature from thrust::square<T>
• catch2_test_device_reduce_nondeterministic.cu: already wraps thrust::square<T>{} internally

---

Checklist

• "New or existing tests cover these changes"
• "The documentation is up to date"

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[rkothari3]

The mypy pre-commit failure seems to be a pre-existing infrastructure issue ( numpy/init.pyi incompatibility with Python 3.14). Changes here are limited to .cu/.cuh files and contain no python. Happy to answer any questions!

[coderabbitai]

📝 Walkthrough

Summary by CodeRabbit

• Refactor

  • Standardized several benchmark operations to use shared squaring utilities instead of local implementations.
  • Updated transform-reduce and for-each benchmark paths to rely on built-in or common in-place squaring behavior.

• Tests

  • Aligned device transform-reduce test coverage with the updated squaring operation used in benchmarks.

• Chores

  • Added a new shared helper for in-place squaring used across multiple benchmark files.

Walkthrough

Adds a shared square_inplace_t<T> device functor to nvbench_helper.cuh and removes locally-defined square_t functors from six benchmark files. for_each/for_each_n benchmarks switch to square_inplace_t, while transform_reduce benchmarks and the CUB device test switch to thrust::square<T>.

Unify square functors

Layer / File(s)	Summary
Add square_inplace_t to nvbench_helper nvbench_helper/nvbench_helper/nvbench_helper.cuh	New square_inplace_t<T> struct template with __device__ void operator()(T& x) providing in-place squaring, available to all benchmarks.
for_each benchmarks adopt square_inplace_t thrust/benchmarks/bench/for_each/basic.cu, thrust/benchmarks/bench/for_each_n/basic.cu, libcudacxx/benchmarks/bench/for_each/basic.cu, libcudacxx/benchmarks/bench/for_each_n/basic.cu	Local square_t device functor definitions removed; operation objects updated to construct square_inplace_t<T> instead.
transform_reduce benchmarks and CUB test adopt thrust::square thrust/benchmarks/bench/transform_reduce/sum.cu, cub/benchmarks/bench/transform_reduce/sum.cu, cub/test/catch2_test_device_transform_reduce.cu	Local square_t definitions removed; thrust/functional.h included; thrust::square<T> used as the transform functor.

Assessment against linked issues

Objective	Addressed	Explanation
Unify square functors by replacing local definitions with thrust::square or a shared alternative [#9619]	✅

Assessment against linked issues: Out-of-scope changes

No out-of-scope changes found. The square_inplace_t addition to nvbench_helper.cuh is necessary because thrust::square is not in-place and cannot directly replace functors used in for_each/for_each_n contexts.

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: c1f1546c-e270-4be2-8a02-9665d1175149

📥 Commits

Reviewing files that changed from the base of the PR and between e1cb696 and 381fdb5.

📒 Files selected for processing (8)

• cub/benchmarks/bench/transform_reduce/sum.cu
• cub/test/catch2_test_device_transform_reduce.cu
• libcudacxx/benchmarks/bench/for_each/basic.cu
• libcudacxx/benchmarks/bench/for_each_n/basic.cu
• nvbench_helper/nvbench_helper/nvbench_helper.cuh
• thrust/benchmarks/bench/for_each/basic.cu
• thrust/benchmarks/bench/for_each_n/basic.cu
• thrust/benchmarks/bench/transform_reduce/sum.cu

[bernhardmgruber]

Great work, thanks!