Repository: NVIDIA/MatX
Pull request #897: Add configurable Pwelch scaling and improve performance
Author: login: tmartin-gh
Created: 2025-03-03T22:53:05Z | Merged: 2025-03-05T18:07:05Z

This PR addresses concerns in issue #887:

* Configurable scaling is supported.  Density and Spectrum modes, similar to Matlab or scipy.signal, as well as fused dB modes frequently needed for spectrum-analyzer-like plotting

* Custom reduction kernel is used instead of MatX sum() using CUB DeviceSegmentedReduce (DSR).

Three reasons why the custom reduction kernel is preferred:

1. In the nominal case where the input signal is a memory backed complex-valued tensor and the output power spectrum is a memory backed real-valued tensor, the intermediate overlapping signal is a 2D tensor {batches, nfft}.  This is an optimal memory layout for the FFTs, but a very suboptimal layout for CUB DSR which would the permuted layout {nfft, batches} as CUB DSR uses one block of threads to read all the 'batches' of a single fft bin.  To improve CUB performance, we'd either need to have the batched FFTs directly write to a permuted tensor (about 4x slower for the scenario considered) or permute the intermediate tensor after the FFT (extra round trip through global memory). 

2. CUB currently needs begin/end indexes for each data segment, the generation of which takes around 50% longer than the custom reduction (ignoring the permutation issue)

3. The custom reduction allows for configurable scaling, such as fusing a dB conversion.


Diff:
diff --git a/docs_input/api/signalimage/general/pwelch.rst b/docs_input/api/signalimage/general/pwelch.rst
index 5e0cb74bd..0223d2ca1 100644
--- a/docs_input/api/signalimage/general/pwelch.rst
+++ b/docs_input/api/signalimage/general/pwelch.rst
@@ -5,8 +5,8 @@ pwelch
 
 Estimate the power spectral density of a signal using Welch's method [1]_
 
-.. doxygenfunction:: pwelch(const xType& x, const wType& w, index_t nperseg, index_t noverlap, index_t nfft)
-.. doxygenfunction:: pwelch(const xType& x, index_t nperseg, index_t noverlap, index_t nfft)
+.. doxygenfunction:: pwelch(const xType& x, const wType& w, index_t nperseg, index_t noverlap, index_t nfft, PwelchOutputScaleMode output_scale_mode, fsType fs)
+.. doxygenfunction:: pwelch(const xType& x, index_t nperseg, index_t noverlap, index_t nfft, PwelchOutputScaleMode output_scale_mode, fsType fs)
 
 Examples
 ~~~~~~~~
diff --git a/examples/pwelch.cu b/examples/pwelch.cu
index f3873900b..46ea34910 100644
--- a/examples/pwelch.cu
+++ b/examples/pwelch.cu
@@ -50,18 +50,15 @@ int main([[maybe_unused]] int argc, [[maybe_unused]] char **argv)
   MATX_ENTER_HANDLER();
   using complex = cuda::std::complex<float>;
 
-  float exec_time_ms;
-  const int num_iterations = 100;
-  index_t signal_size = 256;
-  index_t nperseg = 32;
-  index_t nfft = nperseg;
-  index_t noverlap = 8;
-  float ftone = 3.0;
+  const int num_iterations = 500;
+  index_t signal_size = 256000;
+  index_t nperseg = 512;
+  index_t noverlap = 256;
+  index_t nfft = 65536;
+
+  float ftone = 2048.0;
   cudaStream_t stream;
   cudaStreamCreate(&stream);
-  cudaEvent_t start, stop;
-  cudaEventCreate(&start);
-  cudaEventCreate(&stop);
   cudaExecutor exec{stream};
 
   // Create input signal as a complex exponential
@@ -71,31 +68,30 @@ int main([[maybe_unused]] int argc, [[maybe_unused]] char **argv)
   auto x = make_tensor<complex>({signal_size});
   (x = tmp_x).run(exec); // pre-compute x, tmp_x is otherwise lazily evaluated
 
+  // Create window
+  auto w = make_tensor<float>({nperseg});
+  (w = flattop<0>({nperseg})).run(exec);
+
   // Create output tensor
   auto Pxx  = make_tensor<typename complex::value_type>({nfft});
 
   // Run one time to pre-cache the FFT plan
-  (Pxx = pwelch(x, nperseg, noverlap, nfft)).run(exec);
+  (Pxx = pwelch(x, w, nperseg, noverlap, nfft)).run(exec);
   exec.sync();
 
   // Start the timing
-  cudaEventRecord(start, stream);
-
-  // Start the timing
-  cudaEventRecord(start, stream);
+  exec.start_timer();
 
   for (int iteration = 0; iteration < num_iterations; iteration++) {
     // Use the PWelch operator
-    (Pxx = pwelch(x, nperseg, noverlap, nfft)).run(exec);
+    (Pxx = pwelch(x, w, nperseg, noverlap, nfft)).run(exec);
   }
-
-  cudaEventRecord(stop, stream);
   exec.sync();
-  cudaEventElapsedTime(&exec_time_ms, start, stop);
+  exec.stop_timer();
 
-  printf("Output Pxx:\n");
-  print(Pxx);
-  printf("PWelchOp avg runtime = %.3f ms\n", exec_time_ms / num_iterations);
+  printf("Pxx(0) = %f\n", Pxx(0));
+  printf("Pxx(ftone) = %f\n", Pxx(2048));
+  printf("PWelchOp avg runtime = %.3f ms\n", exec.get_time_ms() / num_iterations);
 
   MATX_CUDA_CHECK_LAST_ERROR();
   MATX_EXIT_HANDLER();
diff --git a/include/matx/kernels/pwelch.cuh b/include/matx/kernels/pwelch.cuh
new file mode 100644
index 000000000..d6af414e6
--- /dev/null
+++ b/include/matx/kernels/pwelch.cuh
@@ -0,0 +1,93 @@
+////////////////////////////////////////////////////////////////////////////////
+// BSD 3-Clause License
+//
+// Copyright (c) 2023, NVIDIA Corporation
+// All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// 1. Redistributions of source code must retain the above copyright notice, this
+//    list of conditions and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice,
+//    this list of conditions and the following disclaimer in the documentation
+//    and/or other materials provided with the distribution.
+//
+// 3. Neither the name of the copyright holder nor the names of its
+//    contributors may be used to endorse or promote products derived from
+//    this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+/////////////////////////////////////////////////////////////////////////////////
+
+#pragma once
+
+namespace matx {
+
+  enum PwelchOutputScaleMode {
+    PwelchOutputScaleMode_Spectrum,
+    PwelchOutputScaleMode_Density,
+    PwelchOutputScaleMode_Spectrum_dB,
+    PwelchOutputScaleMode_Density_dB
+  };
+
+  namespace detail {
+
+#ifdef __CUDACC__
+    template<PwelchOutputScaleMode OUTPUT_SCALE_MODE, typename T_IN, typename T_OUT, typename fsType>
+    __global__ void pwelch_kernel(const T_IN t_in, T_OUT t_out, fsType fs)
+    {
+      const index_t tid = blockIdx.x * blockDim.x + threadIdx.x;
+      const index_t batches = t_in.Shape()[0];
+      const index_t nfft = t_in.Shape()[1];
+
+      if (tid < nfft)
+      {
+        typename T_OUT::value_type pxx = 0;
+        constexpr typename T_OUT::value_type ten = 10;
+
+        for (index_t batch = 0; batch < batches; batch++)
+        {
+          pxx += cuda::std::norm(t_in(batch, tid));
+        }
+
+        if constexpr (OUTPUT_SCALE_MODE == PwelchOutputScaleMode_Spectrum) {
+          t_out(tid) = pxx / batches;
+        }
+        else if constexpr (OUTPUT_SCALE_MODE == PwelchOutputScaleMode_Density) {
+          t_out(tid) = pxx / (batches * fs);
+        }
+        else if constexpr (OUTPUT_SCALE_MODE == PwelchOutputScaleMode_Spectrum_dB) {
+          pxx /= batches;
+          if (pxx != 0) {
+            t_out(tid) = ten * cuda::std::log10(pxx);
+          }
+          else {
+            t_out(tid) = cuda::std::numeric_limits<typename T_OUT::value_type>::lowest();
+          }
+        }
+        else if constexpr (OUTPUT_SCALE_MODE == PwelchOutputScaleMode_Density_dB) {
+          pxx /= (batches * fs);
+          if (pxx != 0) {
+            t_out(tid) = ten * cuda::std::log10(pxx);
+          }
+          else {
+            t_out(tid) = cuda::std::numeric_limits<typename T_OUT::value_type>::lowest();
+          }
+        }
+      }
+    }
+#endif
+
+  };
+};
diff --git a/include/matx/operators/pwelch.h b/include/matx/operators/pwelch.h
index fd1fd1f82..2e1edff94 100644
--- a/include/matx/operators/pwelch.h
+++ b/include/matx/operators/pwelch.h
@@ -40,35 +40,38 @@
 namespace matx
 {
   namespace detail {
-    template <typename OpX, typename OpW>
-    class PWelchOp : public BaseOp<PWelchOp<OpX,OpW>>
+    template <typename OpX, typename OpW, typename fsType>
+    class PWelchOp : public BaseOp<PWelchOp<OpX,OpW,fsType>>
     {
-      private:
-        typename detail::base_type_t<OpX> x_;
-        typename detail::base_type_t<OpW> w_;
-
-        index_t nperseg_;
-        index_t noverlap_;
-        index_t nfft_;
-        cuda::std::array<index_t, 1> out_dims_;
-        mutable detail::tensor_impl_t<typename remove_cvref_t<OpX>::value_type, 1> tmp_out_;
-        mutable typename remove_cvref_t<OpX>::value_type *ptr = nullptr; 
-
       public:
+        static_assert(is_complex_v<typename OpX::value_type>, "pwelch() must have a complex input type");
         using matxop = bool;
         using value_type = typename OpX::value_type::value_type;
         using matx_transform_op = bool;
         using pwelch_xform_op = bool;
 
-        static_assert(is_complex_v<typename OpX::value_type>, "pwelch() must have a complex input type");
 
         __MATX_INLINE__ std::string str() const {
           return "pwelch(" + get_type_str(x_) + "," + get_type_str(w_) + ")";
         }
 
-        __MATX_INLINE__ PWelchOp(const OpX &x, const OpW &w, index_t nperseg, index_t noverlap, index_t nfft) :
-              x_(x), w_(w), nperseg_(nperseg), noverlap_(noverlap), nfft_(nfft) {
-
+        __MATX_INLINE__ PWelchOp(
+              const OpX &x,
+              const OpW &w,
+              index_t nperseg,
+              index_t noverlap,
+              index_t nfft,
+              PwelchOutputScaleMode output_scale_mode,
+              fsType fs
+          ) :
+              x_(x),
+              w_(w),
+              nperseg_(nperseg),
+              noverlap_(noverlap),
+              nfft_(nfft),
+              output_scale_mode_(output_scale_mode),
+              fs_(fs)
+        {
           MATX_STATIC_ASSERT_STR(OpX::Rank() == 1, matxInvalidDim, "pwelch:  Only input rank of 1 is supported presently");
           for (int r = 0; r < OpX::Rank(); r++) {
             out_dims_[r] = nfft_;
@@ -96,7 +99,7 @@ namespace matx
         template <typename Out, typename Executor>
         void Exec(Out &&out, Executor &&ex)  const{
           static_assert(is_cuda_executor_v<Executor>, "pwelch() only supports the CUDA executor currently");
-          pwelch_impl(cuda::std::get<0>(out), x_, w_, nperseg_, noverlap_, nfft_, ex.getStream());
+          pwelch_impl(cuda::std::get<0>(out), x_, w_, nperseg_, noverlap_, nfft_, output_scale_mode_, fs_, ex.getStream());
         }
 
         template <typename ShapeType, typename Executor>
@@ -104,17 +107,17 @@ namespace matx
         {
           if constexpr (is_matx_op<OpX>()) {
             x_.PreRun(std::forward<ShapeType>(shape), std::forward<Executor>(ex));
-          } 
+          }
 
           if constexpr (is_matx_op<OpW>()) {
             w_.PreRun(Shape(w_), std::forward<Executor>(ex));
-          }                 
-        }      
+          }
+        }
 
         template <typename ShapeType, typename Executor>
         __MATX_INLINE__ void PreRun([[maybe_unused]] ShapeType &&shape, Executor &&ex) const noexcept
         {
-          InnerPreRun(std::forward<ShapeType>(shape), std::forward<Executor>(ex));    
+          InnerPreRun(std::forward<ShapeType>(shape), std::forward<Executor>(ex));
 
           detail::AllocateTempTensor(tmp_out_, std::forward<Executor>(ex), out_dims_, &ptr);
 
@@ -133,7 +136,20 @@ namespace matx
           }
 
           matxFree(ptr);
-        }               
+        }
+
+      private:
+        typename detail::base_type_t<OpX> x_;
+        typename detail::base_type_t<OpW> w_;
+
+        index_t nperseg_;
+        index_t noverlap_;
+        index_t nfft_;
+        PwelchOutputScaleMode output_scale_mode_;
+        fsType fs_;
+        cuda::std::array<index_t, 1> out_dims_;
+        mutable detail::tensor_impl_t<typename remove_cvref_t<OpX>::value_type, 1> tmp_out_;
+        mutable typename remove_cvref_t<OpX>::value_type *ptr = nullptr;
     };
   }
 
@@ -144,6 +160,8 @@ namespace matx
    *   Input time domain data type
    * @tparam wType
    *   Input window type
+   * @tparam fsType
+   *   Sampling frequency type
    * @param x
    *   Input time domain tensor
    * @param w
@@ -154,22 +172,69 @@ namespace matx
    *   Number of points to overlap between segments.  Defaults to 0
    * @param nfft
    *   Length of FFT used per segment.  nfft >= nperseg.  Defaults to nfft = nperseg
+   * @param output_scale_mode
+   *   Output scale mode.  Defaults to PwelchOutputScaleMode_Spectrum
+   * @param fs
+   *   Sampling frequency.  Defaults to 1
    *
    * @returns Operator with power spectral density of x
    *
    */
 
-  template <typename xType, typename wType>
-    __MATX_INLINE__ auto pwelch(const xType& x, const wType& w, index_t nperseg, index_t noverlap, index_t nfft)
+  template <
+      typename xType,
+      typename wType,
+      typename fsType>
+    __MATX_INLINE__ auto pwelch(
+        const xType& x,
+        const wType& w,
+        index_t nperseg,
+        index_t noverlap,
+        index_t nfft,
+        PwelchOutputScaleMode output_scale_mode = PwelchOutputScaleMode_Spectrum,
+        fsType fs = 1
+    )
+  {
+    return detail::PWelchOp(x, w, nperseg, noverlap, nfft, output_scale_mode, fs);
+  }
+
+  template <
+      typename xType,
+      typename fsType>
+    __MATX_INLINE__ auto pwelch(
+        const xType& x,
+        index_t nperseg,
+        index_t noverlap,
+        index_t nfft,
+        PwelchOutputScaleMode output_scale_mode = PwelchOutputScaleMode_Spectrum,
+        fsType fs = 1
+    )
   {
-    MATX_NVTX_START("", matx::MATX_NVTX_LOG_API)
+    return detail::PWelchOp(x, std::nullopt, nperseg, noverlap, nfft, output_scale_mode, fs);
+  }
 
-    return detail::PWelchOp(x, w, nperseg, noverlap, nfft);
+  template <typename xType, typename wType>
+    __MATX_INLINE__ auto pwelch(
+        const xType& x,
+        const wType& w,
+        index_t nperseg,
+        index_t noverlap,
+        index_t nfft,
+        PwelchOutputScaleMode output_scale_mode = PwelchOutputScaleMode_Spectrum
+    )
+  {
+    return detail::PWelchOp(x, w, nperseg, noverlap, nfft, output_scale_mode, 1.f);
   }
 
   template <typename xType>
-    __MATX_INLINE__ auto pwelch(const xType& x, index_t nperseg, index_t noverlap, index_t nfft)
+    __MATX_INLINE__ auto pwelch(
+        const xType& x,
+        index_t nperseg,
+        index_t noverlap,
+        index_t nfft,
+        PwelchOutputScaleMode output_scale_mode = PwelchOutputScaleMode_Spectrum
+    )
   {
-    return detail::PWelchOp(x, std::nullopt, nperseg, noverlap, nfft);
+    return detail::PWelchOp(x, std::nullopt, nperseg, noverlap, nfft, output_scale_mode, 1.f);
   }
 }
diff --git a/include/matx/transforms/pwelch.h b/include/matx/transforms/pwelch.h
index 33dfbeed4..c49fbcbbe 100644
--- a/include/matx/transforms/pwelch.h
+++ b/include/matx/transforms/pwelch.h
@@ -32,11 +32,16 @@
 
 #pragma once
 
+#include "matx/kernels/pwelch.cuh"
+
 namespace matx
 {
-  template <typename PxxType, typename xType, typename wType>
-    __MATX_INLINE__ void pwelch_impl(PxxType Pxx, const xType& x, const wType& w, index_t nperseg, index_t noverlap, index_t nfft, cudaStream_t stream=0)
-    {
+  template <typename PxxType, typename xType, typename wType, typename fsType>
+    __MATX_INLINE__ void pwelch_impl(PxxType Pxx, const xType& x, const wType& w, index_t nperseg, index_t noverlap, index_t nfft, PwelchOutputScaleMode output_scale_mode, fsType fs, cudaStream_t stream=0)
+  {
+    #ifndef __CUDACC__
+      MATX_THROW(matxNotSupported, "pwelch not supported on host");
+    #else
       MATX_NVTX_START("", matx::MATX_NVTX_LOG_API)
 
       MATX_ASSERT_STR(Pxx.Rank() == x.Rank(), matxInvalidDim, "pwelch:  Pxx rank must be the same as x rank");
@@ -50,22 +55,28 @@ namespace matx
       index_t batches = x_with_overlaps.Shape()[0];
       auto X_with_overlaps = make_tensor<cuda::std::complex<typename PxxType::value_type>>({batches,static_cast<index_t>(nfft)},MATX_ASYNC_DEVICE_MEMORY,stream);
 
-      if constexpr (std::is_same_v<wType, std::nullopt_t>)
-      {
+      if constexpr (std::is_same_v<wType, std::nullopt_t>) {
         (X_with_overlaps = fft(x_with_overlaps,nfft)).run(stream);
       }
-      else
-      {
+      else {
         (X_with_overlaps = fft(x_with_overlaps * w,nfft)).run(stream);
       }
 
-      // Compute magnitude squared in-place
-      (X_with_overlaps = conj(X_with_overlaps) * X_with_overlaps).run(stream);
-      auto mag_sq_X_with_overlaps = X_with_overlaps.RealView();
-
-      // Perform the reduction across 'batches' rows and normalize
-      auto norm_factor = static_cast<typename PxxType::value_type>(1.) / static_cast<typename PxxType::value_type>(batches);
-      (Pxx = sum(mag_sq_X_with_overlaps, {0}) * norm_factor).run(stream);
-    }
+      int tpb = 512;
+      int bpk = (static_cast<int>(nfft) + tpb - 1) / tpb;
 
+      if (output_scale_mode == PwelchOutputScaleMode_Spectrum) {
+        detail::pwelch_kernel<PwelchOutputScaleMode_Spectrum><<<bpk, tpb, 0, stream>>>(X_with_overlaps, Pxx, fs);
+      }
+      else if (output_scale_mode == PwelchOutputScaleMode_Density) {
+        detail::pwelch_kernel<PwelchOutputScaleMode_Density><<<bpk, tpb, 0, stream>>>(X_with_overlaps, Pxx, fs);
+      }
+      else if (output_scale_mode == PwelchOutputScaleMode_Spectrum_dB) {
+        detail::pwelch_kernel<PwelchOutputScaleMode_Spectrum_dB><<<bpk, tpb, 0, stream>>>(X_with_overlaps, Pxx, fs);
+      }
+      else { //if (output_scale_mode == PwelchOutputScaleMode_Density_dB)
+        detail::pwelch_kernel<PwelchOutputScaleMode_Density_dB><<<bpk, tpb, 0, stream>>>(X_with_overlaps, Pxx, fs);
+      }
+    #endif
+  }
 } // end namespace matx
diff --git a/test/00_operators/PWelch.cu b/test/00_operators/PWelch.cu
index 059bcad13..c43fa90e0 100644
--- a/test/00_operators/PWelch.cu
+++ b/test/00_operators/PWelch.cu
@@ -47,18 +47,23 @@ struct TestParams {
     index_t nperseg;
     index_t noverlap;
     index_t nfft;
+    PwelchOutputScaleMode output_scale_mode;
+    float fs;
     float ftone;
     float sigma;
 };
 
 const std::vector<TestParams> CONFIGS = {
-  {"none", 8, 8, 2, 8, 0., 0.},
-  {"none", 16, 8, 4, 8, 1., 0.},
-  {"none", 16, 8, 4, 8, 2., 1.},
-  {"none", 16384, 256, 64, 256, 63., 0.},
-  {"boxcar", 8, 8, 2, 8, 0., 0.},
-  {"hann", 16, 8, 4, 8, 1., 0.},
-  {"flattop", 1024, 64, 32, 128, 2., 1.},
+  {"none", 8, 8, 2, 8, PwelchOutputScaleMode_Spectrum, 1.0, 0., 0.},
+  {"none", 16, 8, 4, 8, PwelchOutputScaleMode_Spectrum, 1.0, 1., 0.},
+  {"none", 16, 8, 4, 8, PwelchOutputScaleMode_Spectrum, 1.0, 2., 1.},
+  {"none", 16384, 256, 64, 256, PwelchOutputScaleMode_Spectrum, 1.0, 63., 0.},
+  {"boxcar", 8, 8, 2, 8, PwelchOutputScaleMode_Spectrum, 1.0, 0., 0.},
+  {"hann", 16, 8, 4, 8, PwelchOutputScaleMode_Spectrum, 1.0, 1., 0.},
+  {"flattop", 1024, 64, 32, 128, PwelchOutputScaleMode_Spectrum, 2.0, 2., 1.},
+  {"flattop", 1024, 64, 32, 128, PwelchOutputScaleMode_Density, 2.0, 2., 1.},
+  {"flattop", 1024, 64, 32, 128, PwelchOutputScaleMode_Spectrum_dB, 2.0, 2., 1.},
+  {"flattop", 1024, 64, 32, 128, PwelchOutputScaleMode_Density_dB, 2.0, 2., 1.},
 };
 
 class PWelchComplexExponentialTest : public ::testing::TestWithParam<TestParams>
@@ -87,11 +92,26 @@ void helper(PWelchComplexExponentialTest& test)
     "nperseg"_a=test.params.nperseg,
     "noverlap"_a=test.params.noverlap,
     "nfft"_a=test.params.nfft,
+    "scaling"_a="spectrum",
+    "fs"_a =test.params.fs,
     "ftone"_a=test.params.ftone,
     "sigma"_a=test.params.sigma,
     "window_name"_a=test.params.window_name
   );
 
+  if (test.params.output_scale_mode == PwelchOutputScaleMode_Density)
+  {
+    cfg["scaling"] = "density";
+  }
+  else if (test.params.output_scale_mode == PwelchOutputScaleMode_Density_dB)
+  {
+    cfg["scaling"] = "density_dB";
+  }
+  else if (test.params.output_scale_mode == PwelchOutputScaleMode_Spectrum_dB)
+  {
+    cfg["scaling"] = "spectrum_dB";
+  }
+
   test.pb->template InitAndRunTVGeneratorWithCfg<TypeParam>(
       "00_operators", "pwelch_operators", "pwelch_complex_exponential", cfg);
 
@@ -104,7 +124,7 @@ void helper(PWelchComplexExponentialTest& test)
 
   if (test.params.window_name == "none")
   {
-    (Pxx = pwelch(x, test.params.nperseg, test.params.noverlap, test.params.nfft)).run(exec);
+    (Pxx = pwelch(x, test.params.nperseg, test.params.noverlap, test.params.nfft, test.params.output_scale_mode, test.params.fs)).run(exec);
   }
   else
   {
@@ -125,7 +145,7 @@ void helper(PWelchComplexExponentialTest& test)
     {
       ASSERT_TRUE(false) << "Unknown window parameter name " + test.params.window_name;
     }
-    (Pxx = pwelch(x, w, test.params.nperseg, test.params.noverlap, test.params.nfft)).run(exec);
+    (Pxx = pwelch(x, w, test.params.nperseg, test.params.noverlap, test.params.nfft, test.params.output_scale_mode, test.params.fs)).run(exec);
   }
 
   exec.sync();
diff --git a/test/test_vectors/generators/00_operators.py b/test/test_vectors/generators/00_operators.py
index b8948920f..98ed8eee5 100644
--- a/test/test_vectors/generators/00_operators.py
+++ b/test/test_vectors/generators/00_operators.py
@@ -25,24 +25,24 @@ class cross_operator:
     def __init__(self, dtype: str, size: List[int]):
         self.size = size
         self.dtype = dtype
-        np.random.seed(1234)        
+        np.random.seed(1234)
         pass
 
     def run(self) -> Dict[str, np.array]:
         a = matx_common.randn_ndarray(self.size, self.dtype)
         b = matx_common.randn_ndarray(self.size, self.dtype)
-        
+
         #for the result so that the rank matches that of the MatX output
         result = np.cross(a, b)
         if len(result.shape) == 0:
             result = np.expand_dims(result, axis=0)
-            
+
         return {
             'a': a,
             'b': b,
             'out': result
-        }    
-    
+        }
+
 class kron_operator:
     def __init__(self, dtype: str, size: List[int]):
         pass
@@ -130,7 +130,7 @@ def run(self) -> Dict[str, np.array]:
 class toeplitz:
     def __init__(self, dtype: str, size: List[int]):
         self.size = size
-        self.dtype = dtype        
+        self.dtype = dtype
         pass
 
     def run(self) -> Dict[str, np.array]:
@@ -149,7 +149,7 @@ def run(self) -> Dict[str, np.array]:
             'out1': sl.toeplitz(c),
             'out2': sl.toeplitz(c, r),
             'out3': sl.toeplitz(c, r2)
-        }        
+        }
 
 class pwelch_operators:
     def __init__(self, dtype: str, cfg: Dict): #PWelchGeneratorCfg):
@@ -159,6 +159,8 @@ def __init__(self, dtype: str, cfg: Dict): #PWelchGeneratorCfg):
         self.nperseg = cfg['nperseg']
         self.noverlap = cfg['noverlap']
         self.nfft = cfg['nfft']
+        self.scaling = cfg['scaling']
+        self.fs = cfg['fs']
         self.ftone = cfg['ftone']
         self.sigma = cfg['sigma']
         self.window_name = cfg['window_name']
@@ -180,10 +182,16 @@ def pwelch_complex_exponential(self) -> Dict[str, np.ndarray]:
                           noverlap=self.noverlap,
                           nfft=self.nfft,
                           return_onesided=False,
-                          scaling = 'density',
+                          scaling='density',
                           detrend=False)
-        w_scale = np.sum(w*w)
-        Pxx = Pxx * w_scale
+        scale = np.sum(w*w)
+        if self.scaling == 'density' or self.scaling == 'density_dB':
+          scale = scale / self.fs
+        Pxx = Pxx * scale
+
+        if self.scaling == 'density_dB' or self.scaling == 'spectrum_dB':
+          Pxx = 10 * np.log10(Pxx)
+
         return {
             'x_in': x,
             'Pxx_out': Pxx


Reviews:
user: login: cliffburdick
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T00:04:17Z
user: login: cliffburdick
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T05:10:01Z
user: login: cliffburdick
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T05:11:18Z
user: login: cliffburdick
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-04T05:18:20Z
user: login: tmartin-gh
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T17:03:05Z
user: login: tmartin-gh
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T17:03:18Z
user: login: tmartin-gh
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T18:31:09Z
user: login: cliffburdick
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T18:51:39Z
user: login: cliffburdick
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-04T18:53:49Z
user: login: tmartin-gh
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T20:20:15Z

Review comments:
diff_hunk: @@ -71,31 +68,30 @@ int main([[maybe_unused]] int argc, [[maybe_unused]] char **argv)
   auto x = make_tensor<complex>({signal_size});
   (x = tmp_x).run(exec); // pre-compute x, tmp_x is otherwise lazily evaluated
 
+  // Create window
+  auto w = make_tensor<complex>({nperseg});
path: examples/pwelch.cu
user: login: cliffburdick
body: Would this be a tensor usually or just passed as an operator?
created_at: 2025-03-04T00:04:17Z
author_association: COLLABORATOR
diff_hunk: @@ -154,22 +170,41 @@ namespace matx
    *   Number of points to overlap between segments.  Defaults to 0
    * @param nfft
    *   Length of FFT used per segment.  nfft >= nperseg.  Defaults to nfft = nperseg
+   * @param output_scale_mode
+   *   Output scale mode.  Defaults to PwelchOutputScaleMode_Spectrum
+   * @param fs
+   *   Sampling frequency.  Defaults to 1
    *
    * @returns Operator with power spectral density of x
    *
    */
 
   template <typename xType, typename wType>
-    __MATX_INLINE__ auto pwelch(const xType& x, const wType& w, index_t nperseg, index_t noverlap, index_t nfft)
+    __MATX_INLINE__ auto pwelch(
+        const xType& x,
+        const wType& w,
+        index_t nperseg,
+        index_t noverlap,
+        index_t nfft,
+        PwelchOutputScaleMode output_scale_mode = PwelchOutputScaleMode_Spectrum,
+        typename xType::value_type::value_type fs = 1
path: include/matx/operators/pwelch.h
user: login: cliffburdick
body: maybe put a static assert that it's complex so they don't get a weird error if they pass a real?
created_at: 2025-03-04T05:10:00Z
author_association: COLLABORATOR
diff_hunk: @@ -34,8 +34,71 @@
 
 namespace matx
 {
+
path: include/matx/transforms/pwelch.h
user: login: cliffburdick
body: This should go in include/matx/kernels
created_at: 2025-03-04T05:11:18Z
author_association: COLLABORATOR
diff_hunk: @@ -71,31 +68,30 @@ int main([[maybe_unused]] int argc, [[maybe_unused]] char **argv)
   auto x = make_tensor<complex>({signal_size});
   (x = tmp_x).run(exec); // pre-compute x, tmp_x is otherwise lazily evaluated
 
+  // Create window
+  auto w = make_tensor<complex>({nperseg});
path: examples/pwelch.cu
user: login: tmartin-gh
body: Can be tensor or operator
created_at: 2025-03-04T17:03:05Z
author_association: COLLABORATOR
diff_hunk: @@ -34,8 +34,71 @@
 
 namespace matx
 {
+
path: include/matx/transforms/pwelch.h
user: login: tmartin-gh
body: Moved
created_at: 2025-03-04T17:03:18Z
author_association: COLLABORATOR
diff_hunk: @@ -154,22 +170,41 @@ namespace matx
    *   Number of points to overlap between segments.  Defaults to 0
    * @param nfft
    *   Length of FFT used per segment.  nfft >= nperseg.  Defaults to nfft = nperseg
+   * @param output_scale_mode
+   *   Output scale mode.  Defaults to PwelchOutputScaleMode_Spectrum
+   * @param fs
+   *   Sampling frequency.  Defaults to 1
    *
    * @returns Operator with power spectral density of x
    *
    */
 
   template <typename xType, typename wType>
-    __MATX_INLINE__ auto pwelch(const xType& x, const wType& w, index_t nperseg, index_t noverlap, index_t nfft)
+    __MATX_INLINE__ auto pwelch(
+        const xType& x,
+        const wType& w,
+        index_t nperseg,
+        index_t noverlap,
+        index_t nfft,
+        PwelchOutputScaleMode output_scale_mode = PwelchOutputScaleMode_Spectrum,
+        typename xType::value_type::value_type fs = 1
path: include/matx/operators/pwelch.h
user: login: tmartin-gh
body: Done
created_at: 2025-03-04T18:31:09Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,101 @@
+////////////////////////////////////////////////////////////////////////////////
+// BSD 3-Clause License
+//
+// Copyright (c) 2023, NVIDIA Corporation
+// All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// 1. Redistributions of source code must retain the above copyright notice, this
+//    list of conditions and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice,
+//    this list of conditions and the following disclaimer in the documentation
+//    and/or other materials provided with the distribution.
+//
+// 3. Neither the name of the copyright holder nor the names of its
+//    contributors may be used to endorse or promote products derived from
+//    this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+/////////////////////////////////////////////////////////////////////////////////
+
+#pragma once
+
+namespace matx {
+
+  enum PwelchOutputScaleMode {
+    PwelchOutputScaleMode_Spectrum,
+    PwelchOutputScaleMode_Density,
+    PwelchOutputScaleMode_Spectrum_dB,
+    PwelchOutputScaleMode_Density_dB
+  };
+
+  namespace detail {
+
+#ifdef __CUDACC__
+    template<PwelchOutputScaleMode OUTPUT_SCALE_MODE, typename T_IN, typename T_OUT, typename fsType>
+    __global__ void pwelch_kernel(const T_IN t_in, T_OUT t_out, fsType fs)
+    {
+      const index_t tid = blockIdx.x * blockDim.x + threadIdx.x;
+      const index_t batches = t_in.Shape()[0];
+      const index_t nfft = t_in.Shape()[1];
+
+      if (tid < nfft)
+      {
+        typename T_OUT::value_type pxx = 0;
+        constexpr typename T_OUT::value_type ten = 10;
+
+        for (index_t batch = 0; batch < batches; batch++)
+        {
+          pxx += cuda::std::norm(t_in(batch, tid));
+        }
+
+        if constexpr (OUTPUT_SCALE_MODE == PwelchOutputScaleMode_Spectrum)
+        {
+          t_out(tid) = pxx / batches;
+        }
+        else if constexpr (OUTPUT_SCALE_MODE == PwelchOutputScaleMode_Density)
+        {
+          t_out(tid) = pxx / (batches * fs);
+        }
+        else if constexpr (OUTPUT_SCALE_MODE == PwelchOutputScaleMode_Spectrum_dB)
+        {
+          pxx /= batches;
+          if (pxx != 0)
+          {
+            t_out(tid) = ten * cuda::std::log10(pxx);
+          }
+          else
+          {
+            t_out(tid) = cuda::std::numeric_limits<typename T_OUT::value_type>::lowest();
+          }
+        }
+        else if constexpr (OUTPUT_SCALE_MODE == PwelchOutputScaleMode_Density_dB)
+        {
+          pxx /= (batches * fs);
+          if (pxx != 0)
path: include/matx/kernels/pwelch.cuh
user: login: cliffburdick
body: braces on an if/else usually go on the same line to match all the rest of the code. We had a clang format at one point...
created_at: 2025-03-04T18:51:39Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,101 @@
+////////////////////////////////////////////////////////////////////////////////
+// BSD 3-Clause License
+//
+// Copyright (c) 2023, NVIDIA Corporation
+// All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// 1. Redistributions of source code must retain the above copyright notice, this
+//    list of conditions and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice,
+//    this list of conditions and the following disclaimer in the documentation
+//    and/or other materials provided with the distribution.
+//
+// 3. Neither the name of the copyright holder nor the names of its
+//    contributors may be used to endorse or promote products derived from
+//    this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+/////////////////////////////////////////////////////////////////////////////////
+
+#pragma once
+
+namespace matx {
+
+  enum PwelchOutputScaleMode {
+    PwelchOutputScaleMode_Spectrum,
+    PwelchOutputScaleMode_Density,
+    PwelchOutputScaleMode_Spectrum_dB,
+    PwelchOutputScaleMode_Density_dB
+  };
+
+  namespace detail {
+
+#ifdef __CUDACC__
+    template<PwelchOutputScaleMode OUTPUT_SCALE_MODE, typename T_IN, typename T_OUT, typename fsType>
+    __global__ void pwelch_kernel(const T_IN t_in, T_OUT t_out, fsType fs)
+    {
+      const index_t tid = blockIdx.x * blockDim.x + threadIdx.x;
+      const index_t batches = t_in.Shape()[0];
+      const index_t nfft = t_in.Shape()[1];
+
+      if (tid < nfft)
+      {
+        typename T_OUT::value_type pxx = 0;
+        constexpr typename T_OUT::value_type ten = 10;
+
+        for (index_t batch = 0; batch < batches; batch++)
+        {
+          pxx += cuda::std::norm(t_in(batch, tid));
+        }
+
+        if constexpr (OUTPUT_SCALE_MODE == PwelchOutputScaleMode_Spectrum)
+        {
+          t_out(tid) = pxx / batches;
+        }
+        else if constexpr (OUTPUT_SCALE_MODE == PwelchOutputScaleMode_Density)
+        {
+          t_out(tid) = pxx / (batches * fs);
+        }
+        else if constexpr (OUTPUT_SCALE_MODE == PwelchOutputScaleMode_Spectrum_dB)
+        {
+          pxx /= batches;
+          if (pxx != 0)
+          {
+            t_out(tid) = ten * cuda::std::log10(pxx);
+          }
+          else
+          {
+            t_out(tid) = cuda::std::numeric_limits<typename T_OUT::value_type>::lowest();
+          }
+        }
+        else if constexpr (OUTPUT_SCALE_MODE == PwelchOutputScaleMode_Density_dB)
+        {
+          pxx /= (batches * fs);
+          if (pxx != 0)
path: include/matx/kernels/pwelch.cuh
user: login: tmartin-gh
body: Fixed
created_at: 2025-03-04T20:20:14Z
author_association: COLLABORATOR

PR comments:
user: login: tmartin-gh
created_at: 2025-03-03T23:00:19Z
body: /build
author_association: COLLABORATOR
user: login: tmartin-gh
created_at: 2025-03-04T04:04:02Z
body: /build
author_association: COLLABORATOR
user: login: tmartin-gh
created_at: 2025-03-04T17:02:42Z
body: /build
author_association: COLLABORATOR
user: login: tmartin-gh
created_at: 2025-03-04T18:31:19Z
body: /build
author_association: COLLABORATOR
user: login: tmartin-gh
created_at: 2025-03-04T20:20:45Z
body: /build
author_association: COLLABORATOR
user: login: tmartin-gh
created_at: 2025-03-04T23:45:55Z
body: /build
author_association: COLLABORATOR
user: login: tmartin-gh
created_at: 2025-03-04T23:46:03Z
body: /build
author_association: COLLABORATOR
user: login: tmartin-gh
created_at: 2025-03-05T18:07:15Z
body: /build
author_association: COLLABORATOR