Add support for CWT operator

dali/kernels/signal/CMakeLists.txt

@@ -17,6 +17,7 @@ add_subdirectory(decibel)
 if (BUILD_FFTS)
   add_subdirectory(fft)
 endif()
+add_subdirectory(wavelet)
 add_subdirectory(window)
 
 collect_headers(DALI_INST_HDRS PARENT_SCOPE)

dali/kernels/signal/wavelet/CMakeLists.txt

@@ -0,0 +1,17 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+collect_headers(DALI_INST_HDRS PARENT_SCOPE)
+collect_sources(DALI_KERNEL_SRCS PARENT_SCOPE)
+collect_test_sources(DALI_KERNEL_TEST_SRCS PARENT_SCOPE)

dali/kernels/signal/wavelet/cwt_args.h

@@ -0,0 +1,33 @@
+// Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef DALI_KERNELS_SIGNAL_WAVELET_CWT_ARGS_H_
+#define DALI_KERNELS_SIGNAL_WAVELET_CWT_ARGS_H_
+
+namespace dali {
+namespace kernels {
+namespace signal {
+namespace wavelet {
+
+template <typename T = float>
+struct CwtArgs {
+  T a;
+};
+
+}  // namespace wavelet
+}  // namespace signal
+}  // namespace kernels
+}  // namespace dali
+
+#endif  // DALI_KERNELS_SIGNAL_WAVELET_CWT_ARGS_H_

dali/kernels/signal/wavelet/cwt_gpu.cu

@@ -0,0 +1,98 @@
+// Copyright (c) 2020-2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <cmath>
+#include <complex>
+#include <vector>
+#include "dali/core/common.h"
+#include "dali/core/error_handling.h"
+#include "dali/core/format.h"
+#include "dali/kernels/kernel.h"
+#include "dali/kernels/signal/wavelets/cwt_args.h"
+#include "dali/kernels/signal/wavelets/cwt_gpu.h"
+
+namespace dali {
+namespace kernels {
+namespace signal {
+namespace wavelet {
+
+template <typename T>
+struct SampleDesc {
+  const T *in = nullptr;
+  T *out = nullptr;
+  int64_t size = 0;
+};
+
+template <typename T>
+__global__ void CwtKernel(const SampleDesc<T> *sample_data, CwtArgs<T> args) {
+  const int64_t block_size = blockDim.y * blockDim.x;
+  const int64_t grid_size = gridDim.x * block_size;
+  const int sample_idx = blockIdx.y;
+  const auto sample = sample_data[sample_idx];
+  const int64_t offset = block_size * blockIdx.x;
+  const int64_t tid = threadIdx.y * blockDim.x + threadIdx.x;
+
+  for (int64_t idx = offset + tid; idx < sample.size; idx += grid_size) {
+    sample.out[idx] = sample.in[idx] * args.a;
+  }
+}
+
+template <typename T>
+CwtGpu<T>::~CwtGpu() = default;
+
+template <typename T>
+KernelRequirements CwtGpu<T>::Setup(KernelContext &context,
+                                    const InListGPU<T, DynamicDimensions> &in) {
+  auto out_shape = in.shape;
+  const size_t num_samples = in.size();
+  ScratchpadEstimator se;
+  se.add<mm::memory_kind::host, SampleDesc<T>>(num_samples);
+  se.add<mm::memory_kind::device, SampleDesc<T>>(num_samples);
+  KernelRequirements req;
+  req.scratch_sizes = se.sizes;
+  req.output_shapes = {out_shape};
+  return req;
+}
+
+template <typename T>
+void CwtGpu<T>::Run(KernelContext &context, const OutListGPU<T, DynamicDimensions> &out,
+                    const InListGPU<T, DynamicDimensions> &in, const CwtArgs<T> &args) {
+  auto num_samples = in.size();
+  auto *sample_data = context.scratchpad->AllocateHost<SampleDesc<T>>(num_samples);
+
+  for (int i = 0; i < num_samples; i++) {
+    auto &sample = sample_data[i];
+    sample.out = out.tensor_data(i);
+    sample.in = in.tensor_data(i);
+    sample.size = volume(in.tensor_shape(i));
+    assert(sample.size == volume(out.tensor_shape(i)));
+  }
+
+  auto *sample_data_gpu = context.scratchpad->AllocateGPU<SampleDesc<T>>(num_samples);
+  CUDA_CALL(cudaMemcpyAsync(sample_data_gpu, sample_data, num_samples * sizeof(SampleDesc<T>),
+                            cudaMemcpyHostToDevice, context.gpu.stream));
+
+  dim3 block(32, 32);
+  auto blocks_per_sample = std::max(32, 1024 / num_samples);
+  dim3 grid(blocks_per_sample, num_samples);
+  CwtKernel<T><<<grid, block, 0, context.gpu.stream>>>(sample_data_gpu, args);
+}
+
+template class CwtGpu<float>;
+template class CwtGpu<double>;
+
+}  // namespace wavelet
+}  // namespace signal
+}  // namespace kernels
+}  // namespace dali

dali/kernels/signal/wavelet/cwt_gpu.h

@@ -0,0 +1,50 @@
+// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef DALI_KERNELS_SIGNAL_WAVELET_CWT_GPU_H_
+#define DALI_KERNELS_SIGNAL_WAVELET_CWT_GPU_H_
+
+#include <memory>
+#include "dali/core/common.h"
+#include "dali/core/error_handling.h"
+#include "dali/core/format.h"
+#include "dali/core/util.h"
+#include "dali/kernels/kernel.h"
+#include "dali/kernels/signal/wavelets/cwt_args.h"
+
+namespace dali {
+namespace kernels {
+namespace signal {
+namespace wavelet {
+
+template <typename T = float>
+class DLL_PUBLIC CwtGpu {
+ public:
+  static_assert(std::is_floating_point<T>::value, "Only floating point types are supported");
+
+  DLL_PUBLIC ~CwtGpu();
+
+  DLL_PUBLIC KernelRequirements Setup(KernelContext &context,
+                                      const InListGPU<T, DynamicDimensions> &in);
+
+  DLL_PUBLIC void Run(KernelContext &context, const OutListGPU<T, DynamicDimensions> &out,
+                      const InListGPU<T, DynamicDimensions> &in, const CwtArgs<T> &args);
+};
+
+}  // namespace wavelet
+}  // namespace signal
+}  // namespace kernels
+}  // namespace dali
+
+#endif  // DALI_KERNELS_SIGNAL_WAVELET_CWT_GPU_H_

dali/kernels/signal/wavelet/mother_wavelet.cu

@@ -0,0 +1,128 @@
+// Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <cmath>
+#include "dali/kernels/signal/wavelet/mother_wavelet.cuh"
+#include "dali/core/math_util.h"
+
+namespace dali {
+namespace kernels {
+namespace signal {
+
+template <typename T>
+HaarWavelet<T>::HaarWavelet(const std::vector<T> &args) {
+  if (args.size() != 0) {
+    throw new std::invalid_argument("HaarWavelet doesn't accept any arguments.");
+  }
+}
+
+template <typename T>
+__device__ T HaarWavelet<T>::operator()(const T &t) const {
+  if (0.0 <= t && t < 0.5) {
+    return 1.0;
+  }
+  if (0.5 <= t && t < 1.0) {
+    return -1.0;
+  }
+  return 0.0;
+}
+
+template class HaarWavelet<float>;
+template class HaarWavelet<double>;
+
+template <typename T>
+MeyerWavelet<T>::MeyerWavelet(const std::vector<T> &args) {
+  if (args.size() != 0) {
+    throw new std::invalid_argument("MeyerWavelet doesn't accept any arguments.");
+  }
+}
+
+template <typename T>
+__device__ T MeyerWavelet<T>::operator()(const T &t) const {
+  T psi1 = (4/(3*M_PI)*t*std::cos((2*M_PI)/3*t)-1/M_PI*std::sin((4*M_PI)/3*t))/(t-16/9*std::pow(t, 3.0));
+  T psi2 = (8/(3*M_PI)*t*std::cos(8*M_PI/3*t)+1/M_PI*std::sin((4*M_PI)/3)*t)/(t-64/9*std::pow(t, 3.0));
+  return psi1 + psi2;
+}
+
+template class MeyerWavelet<float>;
+template class MeyerWavelet<double>;
+
+template <typename T>
+MexicanHatWavelet<T>::MexicanHatWavelet(const std::vector<T> &args) {
+  if (args.size() != 1) {
+    throw new std::invalid_argument("MexicanHatWavelet accepts exactly one argument - sigma.");
+  }
+  this->sigma = *args.begin();
+}
+
+template <typename T>
+__device__ T MexicanHatWavelet<T>::operator()(const T &t) const {
+  return 2/(std::sqrt(3*sigma)*std::pow(M_PI, 0.25))*(1-std::pow(t/sigma, 2.0))*std::exp(-std::pow(t, 2.0)/(2*std::pow(sigma, 2.0)));
+}
+
+template class MexicanHatWavelet<float>;
+template class MexicanHatWavelet<double>;
+
+template <typename T>
+MorletWavelet<T>::MorletWavelet(const std::vector<T> &args) {
+  if (args.size() != 1) {
+    throw new std::invalid_argument("MorletWavelet accepts exactly 1 argument - C.");
+  }
+  this->C = *args.begin();
+}
+
+template <typename T>
+__device__ T MorletWavelet<T>::operator()(const T &t) const {
+  return C * std::exp(-std::pow(t, 2.0)) * std::cos(5 * t);
+}
+
+template class MorletWavelet<float>;
+template class MorletWavelet<double>;
+
+template <typename T>
+ShannonWavelet<T>::ShannonWavelet(const std::vector<T> &args) {
+  if (args.size() != 0) {
+    throw new std::invalid_argument("ShannonWavelet doesn't accept any arguments.");
+  }
+}
+
+template <typename T>
+__device__ T ShannonWavelet<T>::operator()(const T &t) const {
+  return sinc(t - 0.5) - 2 * sinc(2 * t - 1);
+}
+
+template class ShannonWavelet<float>;
+template class ShannonWavelet<double>;
+
+template <typename T>
+FbspWavelet<T>::FbspWavelet(const std::vector<T> &args) {
+  if (args.size() != 0) {
+    throw new std::invalid_argument("FbspWavelet accepts exactly 3 arguments -> m, fb, fc in that order.");
+  }
+  this->m = *args.begin();
+  this->fb = *(args.begin()+1);
+  this->fc = *(args.begin()+2);
+}
+
+template <typename T>
+__device__ T FbspWavelet<T>::operator()(const T &t) const {
+  return std::sqrt(fb)*std::pow(sinc(t/std::pow(fb, m)), m)*std::exp(2*M_PI*fc*t);
+}
+
+template class FbspWavelet<float>;
+template class FbspWavelet<double>;
+
+}  // namespace signal
+}  // namespace kernel
+}  // namespace dali