add support for functional conv1d

Author: mshr-h

python/tvm/relax/frontend/torch/fx_translator.py

@@ -740,34 +740,54 @@ def _linear_functional(self, node: fx.node.Node) -> relax.Var:
         bias = args[2] if len(args) > 2 else None
         return self.block_builder.emit(relax.op.linear(x, weight, bias, "float32"))
 
-    def _conv1d(self, node: fx.node.Node) -> relax.Var:
-        x = self.env[node.args[0]]
-        module = self.named_modules[node.target]
-        weight = self.params[module.weight]
-
+    def _conv1d_impl(
+        self,
+        x: relax.Expr,
+        weight: relax.Expr,
+        bias: Optional[relax.Expr],
+        strides: Optional[Tuple],
+        padding: Optional[Tuple],
+        dilation: Optional[Tuple],
+        groups: Optional[Tuple],
+    ) -> relax.Var:
         conv1d = self.block_builder.emit(
             relax.op.nn.conv1d(
                 x,
                 weight,
-                strides=module.stride,
-                padding=module.padding,
-                dilation=module.dilation,
-                groups=module.groups,
+                strides=strides,
+                padding=padding,
+                dilation=dilation,
+                groups=groups,
                 data_layout="NCW",
                 kernel_layout="OIW",
                 out_dtype="float32",
             )
         )
 
-        if module.bias is None:
+        if bias is None:
             return conv1d
-
-        bias = self.params[module.bias]
         assert len(self.shape_of(bias)) == 1
         bias = relax.op.reshape(bias, (1, -1, 1))
-
         return self.block_builder.emit(relax.op.add(conv1d, bias))
 
+    def _conv1d(self, node: fx.node.Node) -> relax.Var:
+        x = self.env[node.args[0]]
+        module = self.named_modules[node.target]
+        weight = self.params[module.weight]
+        bias = None
+        if module.bias is not None:
+            bias = self.params[module.bias]
+
+        return self._conv1d_impl(
+            x,
+            weight,
+            bias=bias,
+            strides=module.stride,
+            padding=module.padding,
+            dilation=module.dilation,
+            groups=module.groups,
+        )
+
     def _conv3d(self, node: fx.node.Node) -> relax.Var:
         x = self.env[node.args[0]]
         module = self.named_modules[node.target]
@@ -826,6 +846,25 @@ def _conv2d_impl(
         bias = relax.op.reshape(bias, (1, -1, 1, 1))
         return self.block_builder.emit(relax.op.add(conv2d, bias))
 
+    def _conv1d_functional(self, node: fx.node.Node) -> relax.Var:
+        args = self.retrieve_args(node)
+        x = args[0]
+        weight = args[1]
+        bias = args[2] if len(args) > 2 else None
+        stride = args[3] if len(args) > 3 else 1
+        padding = args[4] if len(args) > 4 else 0
+        dilation = args[5] if len(args) > 5 else 1
+        groups = args[6] if len(args) > 6 else 1
+        return self._conv1d_impl(
+            x,
+            weight,
+            bias=bias,
+            strides=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+        )
+
     def _conv1d_transpose(self, node: fx.node.Node) -> relax.Var:
         x = self.env[node.args[0]]
         module = self.named_modules[node.target]
@@ -1482,6 +1521,7 @@ def create_convert_map(self):
             "type": self._type,
             "astype": self._type,
             "matmul": self._matmul,
+            "conv1d": self._conv1d_functional,
             "conv2d": self._conv2d_functional,
             "linear": self._linear_functional,
             "addmm": self._addmm,