Add op: enable transposed convolution (#4197)

Summary: Pull Request resolved: #4197

Reviewed By: tarun292

Differential Revision: D59589884

fbshipit-source-id: d0b7248b227081e5fefbf4d12cfa2a686d3c646d

Author: manuelcandales

kernels/portable/cpu/op_convolution.cpp

@@ -50,8 +50,10 @@ void conv2d_impl(
     StridesArrayRef out_strides,
     const size_t batch,
     const size_t group,
-    const size_t out_c) {
+    const size_t out_c,
+    bool transposed) {
   size_t in_C = in_sizes[1];
+  size_t out_C = out_sizes[1];
 
   size_t out_H = out_sizes[2];
   size_t in_H = in_sizes[2];
@@ -64,13 +66,15 @@ void conv2d_impl(
   size_t in_C_per_group = in_C / groups;
   size_t in_c_start = group * in_C_per_group;
 
+  size_t out_C_per_group = out_C / groups;
+  size_t out_c_start = group * out_C_per_group;
+
   exec_aten::SizesType in_coord[kTensorDimensionLimit];
   in_coord[0] = batch;
   exec_aten::SizesType out_coord[kTensorDimensionLimit];
   out_coord[0] = batch;
   out_coord[1] = out_c;
   exec_aten::SizesType w_coord[kTensorDimensionLimit];
-  w_coord[0] = out_c;
 
   const int64_t stride_y = val_at(stride, 0);
   const int64_t padding_y = val_at(padding, 0, /*default_value=*/0);
@@ -79,53 +83,115 @@ void conv2d_impl(
   const int64_t padding_x = val_at(padding, 1, /*default_value=*/0);
   const int64_t dilation_x = val_at(dilation, 1);
 
-  // Compute 2D output region
-  for (size_t out_y = 0; out_y < out_H; ++out_y) {
-    out_coord[2] = out_y;
-    for (size_t out_x = 0; out_x < out_W; ++out_x) {
-      out_coord[3] = out_x;
-
-      CTYPE accum = 0.0f;
-      for (size_t in_c = in_c_start; in_c < in_c_start + in_C_per_group;
-           ++in_c) {
-        in_coord[1] = in_c;
-        w_coord[1] = in_c - in_c_start;
-
-        for (size_t w_y = 0; w_y < w_H; ++w_y) {
-          w_coord[2] = w_y;
-
-          size_t in_y = stride_y * out_y + dilation_y * w_y - padding_y;
-          in_coord[2] = in_y;
-          // Only proceed if input y coordinate is within bounds
-          if (in_y >= 0 && in_y < in_H) {
-            for (size_t w_x = 0; w_x < w_W; ++w_x) {
-              w_coord[3] = w_x;
-
-              size_t in_x = stride_x * out_x + dilation_x * w_x - padding_x;
-              in_coord[3] = in_x;
-
-              // Only proceed if input  coordinate is within bounds
-              if (in_x >= 0 && in_x < in_W) {
-                size_t in_idx =
-                    calculate_linear_index(in_coord, in_strides.data(), 4);
-                CTYPE in_val = in_ptr[in_idx];
-
-                size_t w_idx =
-                    calculate_linear_index(w_coord, w_strides.data(), 4);
-                CTYPE w_val = w_ptr[w_idx];
-
-                accum += in_val * w_val;
+  if (!transposed) {
+    w_coord[0] = out_c;
+    // Compute 2D output region
+    for (size_t out_y = 0; out_y < out_H; ++out_y) {
+      out_coord[2] = out_y;
+      for (size_t out_x = 0; out_x < out_W; ++out_x) {
+        out_coord[3] = out_x;
+
+        CTYPE accum = 0.0f;
+        for (size_t in_c = in_c_start; in_c < in_c_start + in_C_per_group;
+             ++in_c) {
+          in_coord[1] = in_c;
+          w_coord[1] = in_c - in_c_start;
+
+          for (size_t w_y = 0; w_y < w_H; ++w_y) {
+            w_coord[2] = w_y;
+
+            size_t in_y = stride_y * out_y + dilation_y * w_y - padding_y;
+            in_coord[2] = in_y;
+            // Only proceed if input y coordinate is within bounds
+            if (in_y >= 0 && in_y < in_H) {
+              for (size_t w_x = 0; w_x < w_W; ++w_x) {
+                w_coord[3] = w_x;
+
+                size_t in_x = stride_x * out_x + dilation_x * w_x - padding_x;
+                in_coord[3] = in_x;
+
+                // Only proceed if input x coordinate is within bounds
+                if (in_x >= 0 && in_x < in_W) {
+                  size_t in_idx =
+                      calculate_linear_index(in_coord, in_strides.data(), 4);
+                  CTYPE in_val = in_ptr[in_idx];
+
+                  size_t w_idx =
+                      calculate_linear_index(w_coord, w_strides.data(), 4);
+                  CTYPE w_val = w_ptr[w_idx];
+
+                  accum += in_val * w_val;
+                }
               }
             }
           }
         }
+
+        if (bias_ptr != nullptr) {
+          accum += convert<CTYPE, CTYPE_BIAS>(bias_ptr[out_c]);
+        }
+        size_t out_idx =
+            calculate_linear_index(out_coord, out_strides.data(), 4);
+        out_ptr[out_idx] = accum;
       }
+    }
+  } else { // transposed convolution
+    if (bias_ptr != nullptr) {
+      out_coord[2] = 0;
+      out_coord[3] = 0;
+      size_t out_c_start_idx =
+          calculate_linear_index(out_coord, out_strides.data(), 4);
+      size_t out_c_end_idx = out_c_start_idx + out_H * out_W;
+      for (size_t out_ix = out_c_start_idx; out_ix < out_c_end_idx; out_ix++) {
+        out_ptr[out_ix] = convert<CTYPE, CTYPE_BIAS>(bias_ptr[out_c]);
+      }
+    }
+
+    w_coord[1] = out_c - out_c_start;
+
+    for (size_t in_y = 0; in_y < in_H; ++in_y) {
+      in_coord[2] = in_y;
+
+      for (size_t in_x = 0; in_x < in_W; ++in_x) {
+        in_coord[3] = in_x;
+
+        for (size_t in_c = in_c_start; in_c < in_c_start + in_C_per_group;
+             ++in_c) {
+          in_coord[1] = in_c;
+
+          size_t in_idx =
+              calculate_linear_index(in_coord, in_strides.data(), 4);
+          CTYPE in_val = in_ptr[in_idx];
+
+          w_coord[0] = in_c;
+          for (size_t w_y = 0; w_y < w_H; ++w_y) {
+            w_coord[2] = w_y;
+            size_t out_y = stride_y * in_y + dilation_y * w_y - padding_y;
+            out_coord[2] = out_y;
+
+            // Only proceed if output y coordinate is within bounds
+            if (out_y >= 0 && out_y < out_H) {
+              for (size_t w_x = 0; w_x < w_W; ++w_x) {
+                w_coord[3] = w_x;
+                size_t out_x = stride_x * in_x + dilation_x * w_x - padding_x;
+                out_coord[3] = out_x;
 
-      if (bias_ptr != nullptr) {
-        accum += convert<CTYPE, CTYPE_BIAS>(bias_ptr[out_c]);
+                // Only proceed if output x coordinate is within bounds
+                if (out_x >= 0 && out_x < out_W) {
+                  size_t w_idx =
+                      calculate_linear_index(w_coord, w_strides.data(), 4);
+                  CTYPE w_val = w_ptr[w_idx];
+
+                  size_t out_idx =
+                      calculate_linear_index(out_coord, out_strides.data(), 4);
+
+                  out_ptr[out_idx] += in_val * w_val;
+                }
+              }
+            }
+          }
+        }
       }
-      size_t out_idx = calculate_linear_index(out_coord, out_strides.data(), 4);
-      out_ptr[out_idx] = accum;
     }
   }
 }
@@ -138,14 +204,9 @@ void convolution_wrapper(
     IntArrayRef stride,
     IntArrayRef padding,
     IntArrayRef dilation,
+    bool transposed,
     int64_t groups,
     Tensor& out) {
-  size_t out_N = in.size(0);
-  size_t out_C = weight.size(0);
-
-  // Compute the number of in and out channels in each group
-  size_t out_C_per_group = out_C / groups;
-
   SizesArrayRef in_sizes = in.sizes();
   SizesArrayRef weight_sizes = weight.sizes();
   SizesArrayRef out_sizes = out.sizes();
@@ -233,6 +294,15 @@ void convolution_wrapper(
   const CTYPE_BIAS* const bias_ptr =
       bias.has_value() ? bias.value().const_data_ptr<CTYPE_BIAS>() : nullptr;
 
+  size_t out_N = out.size(0);
+  size_t out_C_per_group = out.size(1) / groups;
+
+  if (transposed && bias_ptr == nullptr) {
+    // If bias is not present, we need to initialize the output to 0
+    // before we can accumulate into it.
+    memset(out_ptr, 0, out.nbytes());
+  }
+
   for (size_t batch = 0; batch < out_N; ++batch) {
     for (size_t group = 0; group < groups; ++group) {
       // Align channel offset based on the group
@@ -257,7 +327,8 @@ void convolution_wrapper(
             {out_strides, 4},
             batch,
             group,
-            out_c);
+            out_c,
+            transposed);
       }
     }
   }
@@ -273,8 +344,8 @@ Tensor& convolution_out(
     IntArrayRef stride,
     IntArrayRef padding,
     IntArrayRef dilation,
-    __ET_UNUSED bool transposed,
-    __ET_UNUSED IntArrayRef output_padding,
+    bool transposed,
+    IntArrayRef output_padding,
     int64_t groups,
     Tensor& out) {
   (void)ctx;
@@ -298,7 +369,16 @@ Tensor& convolution_out(
   size_t output_ndim = 0;
   exec_aten::SizesType output_sizes[kTensorDimensionLimit];
   get_convolution_out_target_size(
-      in, weight, stride, padding, dilation, output_sizes, &output_ndim);
+      in,
+      weight,
+      stride,
+      padding,
+      dilation,
+      transposed,
+      output_padding,
+      groups,
+      output_sizes,
+      &output_ndim);
 
   ET_KERNEL_CHECK(
       ctx,
@@ -321,12 +401,14 @@ Tensor& convolution_out(
   if (bias.has_value()) {
     bias_type = bias.value().scalar_type();
   }
-  ET_SWITCH_REAL_TYPES(in_type, ctx, "convolution.out", CTYPE, [&]() {
-    ET_SWITCH_REAL_TYPES_AND(
-        Bool, bias_type, ctx, "convolution.out", CTYPE_BIAS, [&]() {
-          convolution_wrapper<CTYPE, CTYPE_BIAS>(
-              in, weight, bias, stride, padding, dilation, groups, out);
-        });
+
+  constexpr auto name = "convolution.out";
+
+  ET_SWITCH_REALH_TYPES(in_type, ctx, name, CTYPE, [&]() {
+    ET_SWITCH_REALHB_TYPES(bias_type, ctx, name, CTYPE_BIAS, [&]() {
+      convolution_wrapper<CTYPE, CTYPE_BIAS>(
+          in, weight, bias, stride, padding, dilation, transposed, groups, out);
+    });
   });
 
   return out;