Skip to content

Commit 480dab7

Browse files
Implement csqrt (#619)
Adding csqrt implementation with casts as c99 csqrt can not be used with nvcc
1 parent 46cc004 commit 480dab7

2 files changed

Lines changed: 18 additions & 0 deletions

File tree

include/matx/operators/scalar_ops.h

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -172,6 +172,23 @@ template <typename T> struct SqrtF {
172172

173173
template <typename T> using SqrtOp = UnOp<T, SqrtF<T>>;
174174

175+
template <typename T>
176+
static __MATX_INLINE__ __MATX_HOST__ __MATX_DEVICE__ auto _internal_csqrt(T v1)
177+
{
178+
static_assert(std::is_floating_point_v<T>, "csqrt() only supports non-complex floating point inputs");
179+
return sqrt(static_cast<cuda::std::complex<T>>(v1));
180+
}
181+
182+
template <typename T> struct CSqrtF {
183+
static __MATX_INLINE__ std::string str() { return "csqrt"; }
184+
static __MATX_INLINE__ __MATX_HOST__ __MATX_DEVICE__ auto op(T v1)
185+
{
186+
return _internal_csqrt(v1);
187+
}
188+
};
189+
190+
191+
template <typename T> using CsqrtOp = UnOp<T, CSqrtF<T>>;
175192

176193
template <typename T>
177194
static __MATX_INLINE__ __MATX_HOST__ __MATX_DEVICE__ auto _internal_conj(T v1)

include/matx/operators/unary_operators.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -362,6 +362,7 @@ namespace matx
362362

363363
#else
364364
DEFINE_UNARY_OP(sqrt, detail::SqrtOp);
365+
DEFINE_UNARY_OP(csqrt, detail::CsqrtOp);
365366
DEFINE_UNARY_OP(exp, detail::ExpOp);
366367
DEFINE_UNARY_OP(expj, detail::ExpjOp);
367368
DEFINE_UNARY_OP(log10, detail::Log10Op);

0 commit comments

Comments
 (0)