add forward portable Conv3d support for contiguous 5D inputs in aten:convolution.out (#19977)

Fixes #11704 

- add forward portable Conv3d support for contiguous 5D inputs in
`aten::convolution.out`
- add direct kernel tests and a portable_lib e2e regression, including
nonzero circular padding


updated according to comments, added tests that cover float16, bfloat16,
int
ci failures seems unrelated to my test

---------

Co-authored-by: RJ Ascani <rja@meta.com>

Author: arkadip-maitra

backends/cortex_m/test/ops/test_conv.py

@@ -333,7 +333,6 @@ def test_dialect_conv2d(test_case, cortex_m_target):
 
 xfails_implementation: dict[str, xfail_type] = {
     "conv1d": "Currently not supported.",
-    "conv3d": "Currently not supported.",
 }
 
 

kernels/portable/cpu/op_convolution.cpp

@@ -191,6 +191,117 @@ void conv2d_impl(
   }
 }
 
+template <typename CTYPE, typename LoadFn = CTYPE (*)(const void*)>
+void conv3d_impl(
+    const CTYPE* const in_ptr,
+    SizesArrayRef in_sizes,
+    StridesArrayRef in_strides,
+    const CTYPE* const w_ptr,
+    SizesArrayRef w_sizes,
+    StridesArrayRef w_strides,
+    const std::optional<Tensor>& bias,
+    const char* const bias_ptr,
+    LoadFn load_bias,
+    IntArrayRef stride,
+    IntArrayRef padding,
+    IntArrayRef dilation,
+    const int64_t groups,
+    CTYPE* const out_ptr,
+    SizesArrayRef out_sizes,
+    StridesArrayRef out_strides,
+    const size_t batch,
+    const size_t group,
+    const size_t out_c) {
+  size_t in_C = in_sizes[1];
+  size_t out_D = out_sizes[2];
+  size_t in_D = in_sizes[2];
+  size_t w_D = w_sizes[2];
+
+  size_t out_H = out_sizes[3];
+  size_t in_H = in_sizes[3];
+  size_t w_H = w_sizes[3];
+
+  size_t out_W = out_sizes[4];
+  size_t in_W = in_sizes[4];
+  size_t w_W = w_sizes[4];
+
+  size_t in_C_per_group = in_C / groups;
+  size_t in_c_start = group * in_C_per_group;
+
+  executorch::aten::SizesType in_coord[kTensorDimensionLimit];
+  in_coord[0] = batch;
+  executorch::aten::SizesType out_coord[kTensorDimensionLimit];
+  out_coord[0] = batch;
+  out_coord[1] = out_c;
+  executorch::aten::SizesType w_coord[kTensorDimensionLimit];
+
+  const int64_t stride_z = val_at(stride, 0);
+  const int64_t padding_z = val_at(padding, 0, /*default_value=*/0);
+  const int64_t dilation_z = val_at(dilation, 0);
+  const int64_t stride_y = val_at(stride, 1);
+  const int64_t padding_y = val_at(padding, 1, /*default_value=*/0);
+  const int64_t dilation_y = val_at(dilation, 1);
+  const int64_t stride_x = val_at(stride, 2);
+  const int64_t padding_x = val_at(padding, 2, /*default_value=*/0);
+  const int64_t dilation_x = val_at(dilation, 2);
+
+  w_coord[0] = out_c;
+  for (const auto out_z : c10::irange(out_D)) {
+    out_coord[2] = out_z;
+    for (const auto out_y : c10::irange(out_H)) {
+      out_coord[3] = out_y;
+      for (const auto out_x : c10::irange(out_W)) {
+        out_coord[4] = out_x;
+
+        CTYPE accum = 0.0f;
+        for (const auto in_c :
+             c10::irange(in_c_start, in_c_start + in_C_per_group)) {
+          in_coord[1] = in_c;
+          w_coord[1] = in_c - in_c_start;
+
+          for (const auto w_z : c10::irange(w_D)) {
+            w_coord[2] = w_z;
+            ssize_t in_z = stride_z * out_z + dilation_z * w_z - padding_z;
+            in_coord[2] = in_z;
+            if (in_z < 0 || in_z >= static_cast<ssize_t>(in_D)) {
+              continue;
+            }
+
+            for (const auto w_y : c10::irange(w_H)) {
+              w_coord[3] = w_y;
+              ssize_t in_y = stride_y * out_y + dilation_y * w_y - padding_y;
+              in_coord[3] = in_y;
+              if (in_y < 0 || in_y >= static_cast<ssize_t>(in_H)) {
+                continue;
+              }
+
+              for (const auto w_x : c10::irange(w_W)) {
+                w_coord[4] = w_x;
+                ssize_t in_x = stride_x * out_x + dilation_x * w_x - padding_x;
+                in_coord[4] = in_x;
+                if (in_x >= 0 && in_x < static_cast<ssize_t>(in_W)) {
+                  size_t in_idx =
+                      calculate_linear_index(in_coord, in_strides.data(), 5);
+                  size_t w_idx =
+                      calculate_linear_index(w_coord, w_strides.data(), 5);
+                  accum += in_ptr[in_idx] * w_ptr[w_idx];
+                }
+              }
+            }
+          }
+        }
+
+        if (bias_ptr != nullptr) {
+          accum += load_bias(&bias_ptr[out_c * bias.value().element_size()]);
+        }
+        size_t out_idx =
+            calculate_linear_index(out_coord, out_strides.data(), 5);
+        out_ptr[out_idx] = accum;
+      }
+    }
+  }
+}
+
 template <typename CTYPE, typename LoadFn = CTYPE (*)(const void*)>
 void convolution_wrapper(
     const Tensor& in,
@@ -294,6 +405,7 @@ void convolution_wrapper(
   size_t out_N = out.size(0);
   size_t out_C = out.size(1);
   size_t out_C_per_group = out_C / groups;
+  const bool is_conv3d = in_sizes.size() == 5;
 
   if (transposed) {
     // For transposed convolution, we need to initialized the output before we
@@ -318,27 +430,50 @@ void convolution_wrapper(
       // Populate all the out channels in the group
       for (const auto out_c :
            c10::irange(out_c_start, out_c_start + out_C_per_group)) {
-        conv2d_impl(
-            in_ptr,
-            in_sizes,
-            {in_strides, 4},
-            w_ptr,
-            weight_sizes,
-            {weight_strides, 4},
-            bias,
-            bias_ptr,
-            load_bias,
-            stride_,
-            padding_,
-            dilation_,
-            groups,
-            out_ptr,
-            out_sizes,
-            {out_strides, 4},
-            batch,
-            group,
-            out_c,
-            transposed);
+        if (is_conv3d) {
+          conv3d_impl(
+              in_ptr,
+              in_sizes,
+              {in_strides, 5},
+              w_ptr,
+              weight_sizes,
+              {weight_strides, 5},
+              bias,
+              bias_ptr,
+              load_bias,
+              stride_,
+              padding_,
+              dilation_,
+              groups,
+              out_ptr,
+              out_sizes,
+              {out_strides, 5},
+              batch,
+              group,
+              out_c);
+        } else {
+          conv2d_impl(
+              in_ptr,
+              in_sizes,
+              {in_strides, 4},
+              w_ptr,
+              weight_sizes,
+              {weight_strides, 4},
+              bias,
+              bias_ptr,
+              load_bias,
+              stride_,
+              padding_,
+              dilation_,
+              groups,
+              out_ptr,
+              out_sizes,
+              {out_strides, 4},
+              batch,
+              group,
+              out_c,
+              transposed);
+        }
       }
     }
   }

kernels/portable/cpu/util/kernel_ops_util.cpp

@@ -45,6 +45,16 @@ bool param_array_is_valid(
   return true;
 }
 
+void fill_convolution_kernel_size(
+    const Tensor& weight,
+    int64_t* kernel_size,
+    size_t* kernel_ndim) {
+  *kernel_ndim = weight.dim() - 2;
+  for (const auto i : c10::irange(*kernel_ndim)) {
+    kernel_size[i] = weight.size(i + 2);
+  }
+}
+
 } // namespace
 
 bool int_array_all_ge(IntArrayRef array, int64_t val) {
@@ -381,18 +391,29 @@ bool check_convolution_args(
     const Tensor& out) {
   ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, weight, out));
 
-  ET_LOG_AND_RETURN_IF_FALSE(tensor_is_default_or_channels_last_dim_order(in));
-  ET_LOG_AND_RETURN_IF_FALSE(
-      tensor_is_default_or_channels_last_dim_order(weight));
-  ET_LOG_AND_RETURN_IF_FALSE(tensor_is_default_or_channels_last_dim_order(out));
-
   ET_CHECK_OR_RETURN_FALSE(
-      in.dim() == 3 || in.dim() == 4,
-      "Expect input tensor to be 3-D or 4-D, but got, %zu.",
+      in.dim() == 3 || in.dim() == 4 || in.dim() == 5,
+      "Expect input tensor to be 3-D, 4-D or 5-D, but got, %zu.",
       static_cast<size_t>(in.dim()));
   ET_LOG_AND_RETURN_IF_FALSE(tensor_is_rank(weight, in.dim()));
   ET_LOG_AND_RETURN_IF_FALSE(tensor_is_rank(out, in.dim()));
 
+  if (in.dim() == 5) {
+    ET_LOG_AND_RETURN_IF_FALSE(tensor_is_default_dim_order(in));
+    ET_LOG_AND_RETURN_IF_FALSE(tensor_is_default_dim_order(weight));
+    ET_LOG_AND_RETURN_IF_FALSE(tensor_is_default_dim_order(out));
+    ET_CHECK_OR_RETURN_FALSE(
+        !transposed,
+        "Transposed 3D convolution is not yet supported on portable.");
+  } else {
+    ET_LOG_AND_RETURN_IF_FALSE(
+        tensor_is_default_or_channels_last_dim_order(in));
+    ET_LOG_AND_RETURN_IF_FALSE(
+        tensor_is_default_or_channels_last_dim_order(weight));
+    ET_LOG_AND_RETURN_IF_FALSE(
+        tensor_is_default_or_channels_last_dim_order(out));
+  }
+
   if (bias.has_value()) {
     ET_LOG_AND_RETURN_IF_FALSE(tensor_is_rank(bias.value(), 1));
     ET_CHECK_OR_RETURN_FALSE(
@@ -404,15 +425,11 @@ bool check_convolution_args(
         transposed ? groups * weight.size(1) : weight.size(0));
   }
 
-  int64_t kernel_size[2];
-  size_t kernel_ndim = 2;
-  if (weight.dim() == 3) {
-    kernel_size[0] = weight.size(2);
-    kernel_ndim = 1;
-  } else {
-    kernel_size[0] = weight.size(2);
-    kernel_size[1] = weight.size(3);
-  }
+  int64_t kernel_size[3];
+  size_t kernel_ndim = 0;
+  fill_convolution_kernel_size(weight, kernel_size, &kernel_ndim);
+  ET_LOG_AND_RETURN_IF_FALSE(
+      kernel_size_is_valid({kernel_size, kernel_ndim}, kernel_ndim));
   ET_LOG_AND_RETURN_IF_FALSE(
       stride_is_valid(stride, kernel_ndim, /*allow_empty=*/false));
   ET_LOG_AND_RETURN_IF_FALSE(
@@ -482,15 +499,9 @@ void get_convolution_out_target_size(
     out_sizes[1] = in.size(1) == 0 ? 0 : groups * weight.size(1);
   }
 
-  int64_t kernel_size[2];
-  size_t kernel_ndim = 2;
-  if (weight.dim() == 3) {
-    kernel_size[0] = weight.size(2);
-    kernel_ndim = 1;
-  } else {
-    kernel_size[0] = weight.size(2);
-    kernel_size[1] = weight.size(3);
-  }
+  int64_t kernel_size[3];
+  size_t kernel_ndim = 0;
+  fill_convolution_kernel_size(weight, kernel_size, &kernel_ndim);
   calculate_kernel_output_sizes(
       in,
       kernel_ndim,