[rocm-libraries] ROCm/rocm-libraries#3836 (commit aef23cd)

[rocThrust] Update complex power function to resolve rounding
 issues (#3836)

## Motivation

Recently, some rounding issues in PyTorch on AMD devices were traced
back to an FMA operation in rocThrusts's pow() function, when type
complex<float> is used.

## Technical Details

This fix (credit to @vpykhtin) adds a specialization for type
complex<float> that is enabled only on AMD GPUs. It performs the
calculation in a different way so that the FMA operation no longer
causes rounding issues.

## Test Plan

Run rocThrust and PyTorch tests. Ensure all tests pass.

## Test Result

rocThrust tests pass locally. Working on running PyTorch tests.

## Submission Checklist

- [ ] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

Author: umfranzw

thrust/detail/complex/cpow.h

@@ -33,6 +33,37 @@ pow(const complex<T0>& x, const complex<T1>& y)
   return exp(log(complex<T>(x)) * complex<T>(y));
 }
 
+// Specialized version for complex<float> on AMD GPUs to use FMA-based multiplication
+#if defined(__HIP_DEVICE_COMPILE__) && defined(__HIP_PLATFORM_AMD__)
+namespace detail
+{
+// FMA-aware complex multiplication for float precision on AMD GPUs.
+// This prevents SLP vectorizer from breaking FMA formation, which causes
+// numerical precision loss in complex arithmetic.
+// The issue occurs when vectorizer packs scalar multiplies before backend
+// can form FMA instructions, resulting in double rounding instead of single.
+THRUST_HOST_DEVICE inline thrust::complex<float> complex_mul_fma(thrust::complex<float> a, thrust::complex<float> b)
+{
+  // Complex multiplication: (a.r + a.i*i) * (b.r + b.i*i)
+  // = (a.r*b.r - a.i*b.i) + (a.r*b.i + a.i*b.r)*i
+  // Using __builtin_fmaf ensures FMA at source level:
+  // real: a.r*b.r + (-(a.i*b.i)) = FMA(a.r, b.r, -(a.i*b.i))
+  // imag: a.i*b.r + a.r*b.i = FMA(a.r, b.i, a.i*b.r)
+  float real_part = __builtin_fmaf(a.real(), b.real(), -(a.imag() * b.imag()));
+  float imag_part = __builtin_fmaf(a.real(), b.imag(), a.imag() * b.real());
+  return thrust::complex<float>(real_part, imag_part);
+}
+} // namespace detail
+
+template <>
+THRUST_HOST_DEVICE inline complex<float> pow(const complex<float>& x, const complex<float>& y)
+{
+  auto log_x   = thrust::log(static_cast<thrust::complex<float>>(x));
+  auto y_log_x = detail::complex_mul_fma(static_cast<thrust::complex<float>>(y), log_x);
+  return static_cast<complex<float>>(thrust::exp(y_log_x));
+}
+#endif
+
 template <typename T0, typename T1>
 THRUST_HOST_DEVICE complex<typename ::internal::promoted_numerical_type<T0, T1>::type>
 pow(const complex<T0>& x, const T1& y)