Add Conv1D operator and layer (#941)

* Add Conv1D operator and layer

* Fix typo in CMakeLists.txt

* Skip the tests when the backend is not available

* Add a separate CMake option for cuDNN

* Minor code cleanup

2 files changed, 143 insertions, 0 deletions

diff --git a/tests/benchmark_ops.cc b/tests/benchmark_ops.cc
index 4d32739c..8ba0fd13 100644
--- a/tests/benchmark_ops.cc
+++ b/tests/benchmark_ops.cc
@@ -105,6 +105,15 @@ void benchmark_dequantize(Device device) {
   BENCHMARK(dequantize_op(x, input_scale, weight_scale, false, true, y), 100000);
 }
 
+void benchmark_conv1d(Device device) {
+  StorageView x({1, 768, 3000}, DataType::FLOAT, device);
+  StorageView weight({768, 768, 3}, DataType::FLOAT, device);
+  StorageView bias({768}, DataType::FLOAT, device);
+  StorageView y(device);
+  const ops::Conv1D conv_op{2, 1};
+  BENCHMARK(conv_op(x, weight, bias, y), 100);
+}
+
 int main(int argc, char* argv[]) {
   if (argc < 3) {
     std::cerr << "usage: " << argv[0] << " op device [dtype]" << std::endl;
@@ -140,6 +149,8 @@ int main(int argc, char* argv[]) {
     benchmark_quantize(device, dtype);
   else if (op == "dequantize")
     benchmark_dequantize(device);
+  else if (op == "conv1d")
+    benchmark_conv1d(device);
 
   return 0;
 }
diff --git a/tests/ops_test.cc b/tests/ops_test.cc
index 9d56a971..a99aac67 100644
--- a/tests/ops_test.cc
+++ b/tests/ops_test.cc
@@ -828,6 +828,138 @@ TEST_P(OpDeviceTest, Max) {
   });
 }
 
+#ifndef CT2_WITH_DNNL
+#  define GUARD_CONV1D_CPU_TEST GTEST_SKIP() << "Conv1D tests on CPU require oneDNN"
+#else
+#  define GUARD_CONV1D_CPU_TEST do {} while (0)
+#endif
+
+#ifndef CT2_WITH_CUDNN
+#  define GUARD_CONV1D_GPU_TEST GTEST_SKIP() << "Conv1D tests on GPU require cuDNN"
+#else
+#  define GUARD_CONV1D_GPU_TEST do {} while (0)
+#endif
+
+static const StorageView conv_input({2, 2, 3}, std::vector<float>{
+    0.5728129f, 0.8784890f, 0.2029965f, 0.3689166f, 0.6570600f, 0.9202735f,
+    0.7081605f, 0.3570334f, 0.9339380f, 0.8162224f, 0.0597404f, 0.4628246f});
+
+static const StorageView conv_weight({4, 2, 2}, std::vector<float>{
+    0.4969918f, 0.3711241f, 0.1489926f, -0.3010672f,
+    -0.2055028f, 0.2540314f, 0.3566069f, -0.1201057f,
+    -0.0737700f, -0.0630847f, -0.2370351f, -0.0451550f,
+    0.0186623f, 0.3600836f, -0.2889268f, -0.4857445f});
+
+static const StorageView conv_bias({4}, std::vector<float>{
+    0.4631361f, -0.1047785f, 0.1047658f, -0.3157263f});
+
+TEST_P(OpDeviceFPTest, Conv1D) {
+  const Device device = GetParam().first;
+  if (device == Device::CUDA)
+    GUARD_CONV1D_GPU_TEST;
+  else
+    GUARD_CONV1D_CPU_TEST;
+  const DataType dtype = GetParam().second;
+  const StorageView expected({2, 4, 2}, std::vector<float>{
+      0.9309945f, 0.7959076f, 0.0533122f, -0.1099610f,
+      -0.1100256f, -0.1701476f, -0.4144599f, -0.8630960f,
+      1.0512151f, 0.8567453f, 0.1242856f, 0.0248157f,
+      -0.1661695f, -0.0155492f, -0.4387956f, -0.2148425f});
+  StorageView output(dtype, device);
+  ops::Conv1D()(conv_input.to(device).to(dtype),
+                conv_weight.to(device).to(dtype),
+                conv_bias.to(device).to(dtype),
+                output);
+  EXPECT_EQ(output.dtype(), dtype);
+  expect_storage_eq(output.to_float(), expected, 1e-3);
+}
+
+TEST_P(OpDeviceFPTest, Conv1DNoBias) {
+  const Device device = GetParam().first;
+  if (device == Device::CUDA)
+    GUARD_CONV1D_GPU_TEST;
+  else
+    GUARD_CONV1D_CPU_TEST;
+  const DataType dtype = GetParam().second;
+  const StorageView expected({2, 4, 2}, std::vector<float>{
+      0.4678584f, 0.3327716f, 0.1580907f, -0.005182412f,
+      -0.2147914f, -0.2749133f, -0.09873369f, -0.5473697f,
+      0.5880789f, 0.3936091f, 0.2290641f, 0.1295942f,
+      -0.2709353f, -0.120315f, -0.1230693f, 0.1008837f});
+  StorageView output(dtype, device);
+  ops::Conv1D()(conv_input.to(device).to(dtype),
+                conv_weight.to(device).to(dtype),
+                output);
+  EXPECT_EQ(output.dtype(), dtype);
+  expect_storage_eq(output.to_float(), expected, 1e-3);
+}
+
+TEST_P(OpDeviceFPTest, Conv1DPadding) {
+  const Device device = GetParam().first;
+  if (device == Device::CUDA)
+    GUARD_CONV1D_GPU_TEST;
+  else
+    GUARD_CONV1D_CPU_TEST;
+  const DataType dtype = GetParam().second;
+  const StorageView expected({2, 4, 4}, std::vector<float>{
+      0.5646521f, 0.9309945f, 0.7959076f, 0.7011377f,
+      -0.0035750f, 0.0533122f, -0.1099610f, 0.1816810f,
+      0.0519716f, -0.1100256f, -0.1701476f, -0.1283464f,
+      -0.2886650f, -0.4144599f, -0.8630960f, -0.5778296f,
+      0.4802138f, 1.0512151f, 0.8567453f, 0.9962531f,
+      -0.0229165f, 0.1242856f, 0.0248157f, -0.1316590f,
+      0.0232352f, -0.1661695f, -0.0155492f, -0.0738365f,
+      -0.4572049f, -0.4387956f, -0.2148425f, -0.4320193f});
+  StorageView output(dtype, device);
+  ops::Conv1D(1, 1)(conv_input.to(device).to(dtype),
+                    conv_weight.to(device).to(dtype),
+                    conv_bias.to(device).to(dtype),
+                    output);
+  EXPECT_EQ(output.dtype(), dtype);
+  expect_storage_eq(output.to_float(), expected, 1e-3);
+}
+
+TEST_P(OpDeviceFPTest, Conv1DStride) {
+  const Device device = GetParam().first;
+  if (device == Device::CUDA)
+    GUARD_CONV1D_GPU_TEST;
+  else
+    GUARD_CONV1D_CPU_TEST;
+  const DataType dtype = GetParam().second;
+  const StorageView expected({2, 4, 1}, std::vector<float>{
+      0.9309945f, 0.0533122f, -0.1100256f, -0.4144599f,
+      1.0512151f, 0.1242856f, -0.1661695f, -0.4387956f});
+  StorageView output(dtype, device);
+  ops::Conv1D(2)(conv_input.to(device).to(dtype),
+                 conv_weight.to(device).to(dtype),
+                 conv_bias.to(device).to(dtype),
+                 output);
+  EXPECT_EQ(output.dtype(), dtype);
+  expect_storage_eq(output.to_float(), expected, 1e-3);
+}
+
+TEST_P(OpDeviceFPTest, Conv1DPaddingAndStride) {
+  const Device device = GetParam().first;
+  if (device == Device::CUDA)
+    GUARD_CONV1D_GPU_TEST;
+  else
+    GUARD_CONV1D_CPU_TEST;
+  const DataType dtype = GetParam().second;
+  const StorageView expected({2, 4, 2}, std::vector<float>{
+      0.5646521f, 0.7959076f, -0.0035750f, -0.1099610f,
+      0.0519716f, -0.1701476f, -0.2886650f, -0.8630960f,
+      0.4802138f, 0.8567453f, -0.0229165f, 0.0248157f,
+      0.0232352f, -0.0155492f, -0.4572049f, -0.2148425f});
+  StorageView output(dtype, device);
+  ops::Conv1D(2, 1)(conv_input.to(device).to(dtype),
+                    conv_weight.to(device).to(dtype),
+                    conv_bias.to(device).to(dtype),
+                    output);
+  EXPECT_EQ(output.dtype(), dtype);
+  expect_storage_eq(output.to_float(), expected, 1e-3);
+}
+
+
 static std::string fp_test_name(::testing::TestParamInfo<std::pair<Device, DataType>> param_info) {
   return dtype_name(param_info.param.second);
 }