Add host direct conv1d and conv2d (#1186)

Implement CPU direct convolution paths for conv1d and conv2d across host executors, matching existing CUDA cropping behavior for full, same, and valid modes.

Add focused host executor tests for 1D and 2D convolution modes and update the executor compatibility documentation.

Tested in localmatx:13.2-arm64: cmake --build build-codex-localmatx --target test_00_transform_ConvCorr -j20; ./test_00_transform_ConvCorr --gtest_filter="CorrelationConvolutionHostDirectTest*"

Author: cliffburdick

docs_input/executor_compatibility.rst

@@ -76,10 +76,10 @@ fused JIT expression; non-JIT CUDA execution through cudaExecutor remains availa
    "clone", "|yes|", "|yes|", "|yes|", "View expression."
    "concat", "|yes|", "|yes|", "|yes|", "View/expression composition."
    "conj", "|yes|", "|yes|", "|yes|", "Element-wise expression."
-   "conv1d", "|partial|", "|yes|", "|no|", "Direct CUDA convolution is CUDA-only; FFT convolution can use the CPU FFT backend when enabled."
-   "conv2d", "|no|", "|yes|", "|no|", "CUDA-only convolution transform."
+   "conv1d", "|yes|", "|yes|", "|no|", "Direct convolution supports host and CUDA executors. FFT convolution can use the CPU FFT backend when enabled."
+   "conv2d", "|yes|", "|yes|", "|no|", "Direct convolution supports host and CUDA executors."
    "copy", "|yes|", "|yes|", "|no|", "Executor-dispatched assignment/copy transform; CUDAJITExecutor fuses expression evaluation instead of using this transform directly."
-   "corr", "|partial|", "|yes|", "|no|", "Correlation transform; host FFT backend may be required depending on mode and input."
+   "corr", "|yes|", "|yes|", "|no|", "Direct correlation supports host and CUDA executors. FFT correlation can use the CPU FFT backend when enabled."
    "cos", "|yes|", "|yes|", "|yes|", "Element-wise expression."
    "cosh", "|yes|", "|yes|", "|yes|", "Element-wise expression."
    "cov", "|no|", "|yes|", "|no|", "CUDA-only covariance transform."

include/matx/executors/support.h

@@ -83,12 +83,7 @@ constexpr bool CheckFFTSupport() {
 
 template <typename Exec>
 constexpr bool CheckDirect1DConvSupport() {
-  if constexpr (is_host_executor_v<Exec>) {
-    return false;
-  }
-  else {
-    return true;
-  }
+  return true;
 }
 
 template <typename Exec, typename T>
@@ -103,12 +98,7 @@ constexpr bool CheckFFT1DConvSupport() {
 
 template <typename Exec>
 constexpr bool Check2DConvSupport() {
-  if constexpr (is_host_executor_v<Exec>) {
-    return false;
-  }
-  else {
-    return true;
-  }
+  return true;
 }
 
 template <typename Exec, typename T>

include/matx/operators/conv.h

@@ -177,7 +177,6 @@ namespace matx
 
         template <typename Out, typename Executor>
         void Exec(Out &&out, Executor &&ex) const {
-          MATX_ASSERT_STR(!(is_host_executor_v<Executor> && method_ == MATX_C_METHOD_DIRECT), matxNotSupported, "direct conv1d() only supports the CUDA executor currently");
           MATX_STATIC_ASSERT_STR((Rank() == cuda::std::tuple_element_t<0, remove_cvref_t<Out>>::Rank()), 
                 matxInvalidParameter, "conv1d: inputs and outputs must have same rank to use conv1d with axis parameter");
           if constexpr (!std::is_same_v<PermDims, no_permute_t>) {
@@ -405,13 +404,11 @@ namespace detail {
 
       template <typename Out, typename Executor>
       void Exec(Out &&out, Executor &&ex) const {
-        static_assert(is_cuda_executor_v<Executor>, "conv2d() only supports the CUDA executor currently");
-
         if constexpr (!std::is_same_v<PermDims, no_permute_t>) {
-          conv2d_impl(permute(cuda::std::get<0>(out), perm_), a_, b_, mode_, ex.getStream());
+          conv2d_impl(permute(cuda::std::get<0>(out), perm_), a_, b_, mode_, ex);
         }
         else {
-          conv2d_impl(cuda::std::get<0>(out), a_, b_, mode_, ex.getStream());
+          conv2d_impl(cuda::std::get<0>(out), a_, b_, mode_, ex);
         }
       }
 

include/matx/operators/corr.h

@@ -158,7 +158,6 @@ namespace matx
 
         template <typename Out, typename Executor>
         void Exec(Out &&out, Executor &&ex) const {
-          MATX_ASSERT_STR(!(is_host_executor_v<Executor> && method_ == MATX_C_METHOD_DIRECT), matxNotSupported, "direct corr() only supports the CUDA executor currently");
           MATX_STATIC_ASSERT_STR((Rank() == cuda::std::tuple_element_t<0, remove_cvref_t<Out>>::Rank()), 
                 matxInvalidParameter, "corr: inputs and outputs must have same rank to use corr with axis parameter");
           if constexpr (!std::is_same_v<PermDims, no_permute_t>) {
@@ -253,4 +252,4 @@ __MATX_INLINE__ auto corr(const In1Type &i1, const In2Type &i2,
   return detail::CorrOp(in1, in2, mode, method, perm);
 }
 
-}
+}

include/matx/transforms/conv.h

@@ -223,6 +223,71 @@ inline void matxDirectConv1DInternal(OutputType &o, const InType &i,
 #endif
 }
 
+template <typename OutputType, typename InType, typename FilterType, typename Executor>
+  requires is_host_executor_v<Executor>
+inline void matxDirectConv1DInternal(OutputType &o, const InType &i,
+                                     const FilterType &filter, matxConvCorrMode_t mode,
+                                     [[maybe_unused]] const Executor &exec)
+{
+  MATX_NVTX_START("", matx::MATX_NVTX_LOG_INTERNAL)
+
+  MATX_STATIC_ASSERT(OutputType::Rank() == InType::Rank(), matxInvalidDim);
+  MATX_STATIC_ASSERT(OutputType::Rank() == FilterType::Rank(), matxInvalidDim);
+
+  constexpr int Rank = OutputType::Rank();
+  const index_t signal_len = i.Size(Rank - 1);
+  const index_t filter_len = filter.Size(Rank - 1);
+  [[maybe_unused]] const index_t full_len = signal_len + filter_len - 1;
+
+  MATX_ASSERT_STR(mode != MATX_C_MODE_FULL || o.Size(Rank - 1) == full_len,
+      matxInvalidSize, "Output size for FULL convolution incorrect");
+  MATX_ASSERT_STR(mode != MATX_C_MODE_SAME || o.Size(Rank - 1) == signal_len,
+      matxInvalidSize, "Output size for SAME convolution incorrect");
+  MATX_ASSERT_STR(mode != MATX_C_MODE_VALID || o.Size(Rank - 1) == signal_len - filter_len + 1,
+      matxInvalidSize, "Output size for VALID convolution incorrect");
+
+  index_t start = 0;
+  if (mode == MATX_C_MODE_SAME) {
+    start = (filter_len & 1) ? (filter_len - 1) / 2 : filter_len / 2 - 1;
+  }
+  else if (mode == MATX_C_MODE_VALID) {
+    start = filter_len - 1;
+  }
+
+  const index_t total = static_cast<index_t>(TotalSize(o));
+
+#ifdef MATX_EN_OMP
+#pragma omp parallel for num_threads(exec.GetNumThreads()) if(exec.GetNumThreads() > 1)
+#endif
+  for (index_t abs = 0; abs < total; abs++) {
+    auto odims = GetIdxFromAbs(o, abs);
+    const index_t full_idx = odims[Rank - 1] + start;
+
+    typename OutputType::value_type sum{0};
+    for (index_t f = 0; f < filter_len; f++) {
+      const index_t sig_idx = full_idx - f;
+      if (sig_idx >= 0 && sig_idx < signal_len) {
+        auto idims = odims;
+        auto fdims = odims;
+        idims[Rank - 1] = sig_idx;
+        fdims[Rank - 1] = f;
+
+        const auto ival = cuda::std::apply([&](auto &&...args) {
+          return i.operator()(args...);
+        }, idims);
+        const auto fval = cuda::std::apply([&](auto &&...args) {
+          return filter.operator()(args...);
+        }, fdims);
+        sum = detail::madd(ival, fval, sum);
+      }
+    }
+
+    cuda::std::apply([&](auto &&...args) {
+      o.operator()(args...) = sum;
+    }, odims);
+  }
+}
+
 
 template <typename OutputType, typename In1Type, typename In2Type>
 void matxDirectConv2DInternal(OutputType &o, In1Type &in1,
@@ -256,6 +321,82 @@ void matxDirectConv2DInternal(OutputType &o, In1Type &in1,
   Conv2D<OutputType, In1Type, In2Type, BLOCK_X, BLOCK_Y, FILTER_SHARED_X, FILTER_SHARED_Y, FILTER_REG_X, FILTER_REG_Y, ILPY><<<blocks, threads, 0, stream>>>(o, in1, in2, mode, num_batch);
 #endif
 }
+
+template <typename OutputType, typename In1Type, typename In2Type, typename Executor>
+  requires is_host_executor_v<Executor>
+void matxDirectConv2DInternal(OutputType &o, const In1Type &in1,
+                              const In2Type &in2, matxConvCorrMode_t mode,
+                              [[maybe_unused]] const Executor &exec)
+{
+  MATX_NVTX_START("", matx::MATX_NVTX_LOG_INTERNAL)
+
+  MATX_STATIC_ASSERT(OutputType::Rank() == In1Type::Rank(), matxInvalidDim);
+  MATX_STATIC_ASSERT(OutputType::Rank() == In2Type::Rank(), matxInvalidDim);
+  MATX_STATIC_ASSERT(OutputType::Rank() >= 2, matxInvalidDim);
+
+  constexpr int Rank = OutputType::Rank();
+
+  const index_t i1N = in1.Size(Rank - 2);
+  const index_t i1M = in1.Size(Rank - 1);
+  const index_t i2N = in2.Size(Rank - 2);
+  const index_t i2M = in2.Size(Rank - 1);
+
+  index_t dy = 0;
+  index_t dx = 0;
+  if (mode == MATX_C_MODE_SAME) {
+    dy = i2N / 2;
+    dx = i2M / 2;
+  }
+  else if (mode == MATX_C_MODE_FULL) {
+    dy = i2N - 1;
+    dx = i2M - 1;
+  }
+
+  const index_t total = static_cast<index_t>(TotalSize(o));
+
+#ifdef MATX_EN_OMP
+#pragma omp parallel for num_threads(exec.GetNumThreads()) if(exec.GetNumThreads() > 1)
+#endif
+  for (index_t abs = 0; abs < total; abs++) {
+    auto odims = GetIdxFromAbs(o, abs);
+    const index_t oy = odims[Rank - 2];
+    const index_t ox = odims[Rank - 1];
+
+    typename OutputType::value_type sum{0};
+    for (index_t fy = 0; fy < i2N; fy++) {
+      const index_t iy = oy + fy - dy;
+      if (iy < 0 || iy >= i1N) {
+        continue;
+      }
+
+      for (index_t fx = 0; fx < i2M; fx++) {
+        const index_t ix = ox + fx - dx;
+        if (ix < 0 || ix >= i1M) {
+          continue;
+        }
+
+        auto i1dims = odims;
+        auto i2dims = odims;
+        i1dims[Rank - 2] = iy;
+        i1dims[Rank - 1] = ix;
+        i2dims[Rank - 2] = i2N - 1 - fy;
+        i2dims[Rank - 1] = i2M - 1 - fx;
+
+        const auto i1val = cuda::std::apply([&](auto &&...args) {
+          return in1.operator()(args...);
+        }, i1dims);
+        const auto i2val = cuda::std::apply([&](auto &&...args) {
+          return in2.operator()(args...);
+        }, i2dims);
+        sum = detail::madd(i1val, i2val, sum);
+      }
+    }
+
+    cuda::std::apply([&](auto &&...args) {
+      o.operator()(args...) = sum;
+    }, odims);
+  }
+}
 } // end namespace detail
 
 template <typename OutputType, typename In1Type, typename In2Type, typename Executor>
@@ -401,42 +542,55 @@ inline void conv1d_impl(OutputType o, const In1Type &i1, const In2Type &i2,
  * @param in1 First input operator
  * @param in2 Second input operator
  * @param mode Convolution mode
- * @param stream CUDA stream
+ * @param exec Executor
  */
-template <typename OutputType, typename In1Type, typename In2Type>
+template <typename OutputType, typename In1Type, typename In2Type, typename Executor>
 inline void conv2d_impl(OutputType o, const In1Type in1, const In2Type in2,
-                   matxConvCorrMode_t mode, cudaStream_t stream = 0)
+                   matxConvCorrMode_t mode, const Executor &exec)
 {
   MATX_NVTX_START("", matx::MATX_NVTX_LOG_API)
   constexpr int Rank1 = In1Type::Rank();
   constexpr int Rank2 = In2Type::Rank();
 
   if constexpr (In1Type::Rank() == In2Type::Rank()) {
-     index_t size1 = in1.Size(Rank1-1) * in1.Size(Rank1-2);
-     index_t size2 = in2.Size(Rank2-1) * in2.Size(Rank2-2);
-     // smaller size is the filter, set it as second input
-     if(size1 >= size2) {
-       detail::matxDirectConv2DInternal(o, in1, in2, mode, stream);
-     } else {  // swap in1/in2
-       detail::matxDirectConv2DInternal(o, in2, in1, mode, stream);
-     }
-  // These branches clone the inputs to match in rank
-  } else if constexpr (In1Type::Rank() <In2Type::Rank()) {
-      // in1 is smaller so clone it to match in2
-      auto shape = in2.Shape();
-      int d = Rank2 - Rank1;
-      for(int i = 0; i < Rank1; i++) {
-        shape[i+d] = matxKeepDim;
+    auto run_direct = [&](const auto &signal, const auto &filter) {
+      if constexpr (is_cuda_executor_v<Executor>) {
+        detail::matxDirectConv2DInternal(o, signal, filter, mode, exec.getStream());
       }
-      conv2d_impl(o, clone<Rank2>(in1, shape), in2, mode, stream);
-  } else {
-      // in1 is smaller so clone it to match in2
-      auto shape = in1.Shape();
-      int d = Rank1 - Rank2;
-      for(int i = 0; i < Rank2; i++) {
-        shape[i+d] = matxKeepDim;
+      else if constexpr (is_host_executor_v<Executor>) {
+        detail::matxDirectConv2DInternal(o, signal, filter, mode, exec);
+      }
+      else {
+        static_assert(is_cuda_executor_v<Executor> || is_host_executor_v<Executor>,
+                      "conv2d() only supports CUDA and host executors currently");
       }
-      conv2d_impl(o, in1, clone<Rank1>(in2, shape), mode, stream);
+    };
+
+    index_t size1 = in1.Size(Rank1-1) * in1.Size(Rank1-2);
+    index_t size2 = in2.Size(Rank2-1) * in2.Size(Rank2-2);
+    // smaller size is the filter, set it as second input
+    if(size1 >= size2) {
+      run_direct(in1, in2);
+    } else {  // swap in1/in2
+      run_direct(in2, in1);
+    }
+  // These branches clone the inputs to match in rank
+  } else if constexpr (In1Type::Rank() < In2Type::Rank()) {
+    // in1 is smaller so clone it to match in2
+    auto shape = in2.Shape();
+    int d = Rank2 - Rank1;
+    for(int i = 0; i < Rank1; i++) {
+      shape[i+d] = matxKeepDim;
+    }
+    conv2d_impl(o, clone<Rank2>(in1, shape), in2, mode, exec);
+  } else {
+    // in1 is smaller so clone it to match in2
+    auto shape = in1.Shape();
+    int d = Rank1 - Rank2;
+    for(int i = 0; i < Rank2; i++) {
+      shape[i+d] = matxKeepDim;
+    }
+    conv2d_impl(o, in1, clone<Rank1>(in2, shape), mode, exec);
   }
 }