[Relax][PyTorch] Add AvgPool 1D and 3D Op Support for Exported Program and FX graph

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

@@ -580,6 +580,48 @@ def _addmm(self, node: fx.Node) -> relax.Var:
             res = bias if res is None else self.block_builder.emit(relax.op.add(bias, res))
         return res
 
+    def _avg_pool1d_impl(
+        self,
+        x: relax.Expr,
+        kernel_size: Union[int, Tuple[int]] = 1,
+        stride: Optional[Union[int, Tuple[int]]] = None,
+        padding: Optional[int] = 0,
+        ceil_mode: Optional[bool] = False,
+        count_include_pad: Optional[bool] = True,
+    ) -> relax.Var:
+        # Expand to 3D by adding batch dim if input is 2D
+        x_ndim = x.struct_info.ndim
+        if x_ndim == 2:
+            x = relax.op.expand_dims(x, axis=0)
+        stride = kernel_size if stride is None or stride == [] else stride
+
+        result = self.block_builder.emit(
+            relax.op.nn.avg_pool1d(
+                x,
+                pool_size=kernel_size,
+                strides=stride,
+                padding=padding,
+                ceil_mode=ceil_mode,
+                count_include_pad=count_include_pad,
+                layout="NCW",
+            )
+        )
+        # Remove added batch dim from result
+        if x_ndim == 2:
+            result = relax.op.squeeze(result, axis=[0])
+        return result
+
+    def _avg_pool1d(self, node: fx.Node) -> relax.Var:
+        args, kwargs = node.normalized_arguments(node)
+        x = self.env[args[0]]
+        kernel_size = args[1] if len(args) > 1 else kwargs["kernel_size"]
+        stride = args[2] if len(args) > 2 else kwargs.get("stride", None)
+        padding = args[3] if len(args) > 3 else kwargs.get("padding", 0)
+        ceil_mode = args[4] if len(args) > 4 else kwargs.get("ceil_mode", False)
+        count_include_pad = args[5] if len(args) > 5 else kwargs.get("count_include_pad", True)
+
+        return self._avg_pool1d_impl(x, kernel_size, stride, padding, ceil_mode, count_include_pad)
+
     def _avg_pool2d_impl(
         self,
         x: relax.Expr,
@@ -588,8 +630,13 @@ def _avg_pool2d_impl(
         padding: Optional[int] = 0,
         ceil_mode: Optional[bool] = False,
     ) -> relax.Var:
+        # Expand to 4D by adding batch dim if input is 3D
+        x_ndim = x.struct_info.ndim
+        if x_ndim == 3:
+            x = relax.op.expand_dims(x, axis=0)
         stride = kernel_size if stride is None or stride == [] else stride
-        return self.block_builder.emit(
+
+        result = self.block_builder.emit(
             relax.op.nn.avg_pool2d(
                 x,
                 pool_size=kernel_size,
@@ -599,6 +646,10 @@ def _avg_pool2d_impl(
                 layout="NCHW",
             )
         )
+        # Remove added batch dim from result
+        if x_ndim == 3:
+            result = relax.op.squeeze(result, axis=[0])
+        return result
 
     def _avg_pool2d(self, node: fx.Node) -> relax.Var:
         args, kwargs = node.normalized_arguments(node)
@@ -609,6 +660,48 @@ def _avg_pool2d(self, node: fx.Node) -> relax.Var:
         ceil_mode = args[4] if len(args) > 4 else kwargs.get("ceil_mode", False)
         return self._avg_pool2d_impl(x, kernel_size, stride, padding, ceil_mode)
 
+    def _avg_pool3d_impl(
+        self,
+        x: relax.Expr,
+        kernel_size: Union[int, Tuple[int, int, int]] = (1, 1, 1),
+        stride: Optional[Union[int, Tuple[int, int, int]]] = None,
+        padding: Optional[int] = 0,
+        ceil_mode: Optional[bool] = False,
+        count_include_pad: Optional[bool] = True,
+    ) -> relax.Var:
+        # Expand to 5D by adding batch dim if input is 4D
+        x_ndim = x.struct_info.ndim
+        if x_ndim == 4:
+            x = relax.op.expand_dims(x, axis=0)
+        stride = kernel_size if stride is None or stride == [] else stride
+
+        result = self.block_builder.emit(
+            relax.op.nn.avg_pool3d(
+                x,
+                pool_size=kernel_size,
+                strides=stride,
+                padding=padding,
+                ceil_mode=ceil_mode,
+                count_include_pad=count_include_pad,
+                layout="NCDHW",
+            )
+        )
+        # Remove added batch dim from result
+        if x_ndim == 4:
+            result = relax.op.squeeze(result, axis=[0])
+        return result
+
+    def _avg_pool3d(self, node: fx.Node) -> relax.Var:
+        args, kwargs = node.normalized_arguments(node)
+        x = self.env[args[0]]
+        kernel_size = args[1] if len(args) > 1 else kwargs["kernel_size"]
+        stride = args[2] if len(args) > 2 else kwargs.get("stride", None)
+        padding = args[3] if len(args) > 3 else kwargs.get("padding", 0)
+        ceil_mode = args[4] if len(args) > 4 else kwargs.get("ceil_mode", False)
+        count_include_pad = args[5] if len(args) > 5 else kwargs.get("count_include_pad", True)
+
+        return self._avg_pool3d_impl(x, kernel_size, stride, padding, ceil_mode, count_include_pad)
+
     def _baddbmm(self, node: fx.Node) -> relax.Var:
         x = self.env[node.args[0]]
         batch1 = self.env[node.args[1]]

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

@@ -401,7 +401,9 @@ def create_convert_map(
             "adaptive_avg_pool2d.default": self._adaptive_avg_pool2d,
             "adaptive_avg_pool3d.default": self._adaptive_avg_pool3d,
             "addmm.default": self._addmm,
+            "avg_pool1d.default": self._avg_pool1d,
             "avg_pool2d.default": self._avg_pool2d,
+            "avg_pool3d.default": self._avg_pool3d,
             "baddbmm.default": self._baddbmm,
             "bmm.default": self._binary_op(
                 partial(relax.op.linear_algebra.matmul, out_dtype="float32"), operator.matmul

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

@@ -230,6 +230,15 @@ def _adaptive_avg_pool3d_module(self, node: fx.Node) -> relax.Var:
             result = relax.op.squeeze(result, axis=[0])
         return result
 
+    def _avg_pool1d_module(self, node: fx.Node) -> relax.Var:
+        x = self.env[node.args[0]]
+        module = self.named_modules[node.target]
+        kernel_size = module.kernel_size
+        stride = module.stride
+        padding = module.padding
+        ceil_mode = module.ceil_mode
+        return self._avg_pool1d_impl(x, kernel_size, stride, padding, ceil_mode)
+
     def _avg_pool2d_module(self, node: fx.Node) -> relax.Var:
         x = self.env[node.args[0]]
         module = self.named_modules[node.target]
@@ -239,6 +248,15 @@ def _avg_pool2d_module(self, node: fx.Node) -> relax.Var:
         ceil_mode = module.ceil_mode
         return self._avg_pool2d_impl(x, kernel_size, stride, padding, ceil_mode)
 
+    def _avg_pool3d_module(self, node: fx.Node) -> relax.Var:
+        x = self.env[node.args[0]]
+        module = self.named_modules[node.target]
+        kernel_size = module.kernel_size
+        stride = module.stride
+        padding = module.padding
+        ceil_mode = module.ceil_mode
+        return self._avg_pool3d_impl(x, kernel_size, stride, padding, ceil_mode)
+
     def _batch_norm_2d_module(self, node: fx.Node) -> relax.Var:
         x = self.env[node.args[0]]
         module = self.named_modules[node.target]
@@ -711,7 +729,9 @@ def create_convert_map(
             nn.AdaptiveAvgPool1d: self._adaptive_avg_pool1d_module,
             nn.AdaptiveAvgPool2d: self._adaptive_avg_pool2d_module,
             nn.AdaptiveAvgPool3d: self._adaptive_avg_pool3d_module,
+            nn.AvgPool1d: self._avg_pool1d_module,
             nn.AvgPool2d: self._avg_pool2d_module,
+            nn.AvgPool3d: self._avg_pool3d_module,
             nn.BatchNorm2d: self._batch_norm_2d_module,
             nn.Conv1d: self._conv1d_module,
             nn.Conv2d: self._conv2d_module,
@@ -824,7 +844,9 @@ def create_convert_map(
             "adaptive_avg_pool2d": self._adaptive_avg_pool2d,
             "adaptive_avg_pool3d": self._adaptive_avg_pool3d,
             "addmm": self._addmm,
+            "avg_pool1d": self._avg_pool1d,
             "avg_pool2d": self._avg_pool2d,
+            "avg_pool3d": self._avg_pool3d,
             "baddbmm": self._baddbmm,
             "bmm": self._binary_op(
                 partial(relax.op.linear_algebra.matmul, out_dtype="float32"), operator.matmul

python/tvm/relax/op/nn/nn.py

@@ -840,7 +840,7 @@ def avg_pool1d(
     padding: Union[int, Tuple[int, ...]] = (0, 0),
     dilation: Union[int, Tuple[int, int]] = (1,),
     ceil_mode: bool = False,
-    count_include_pad: bool = False,
+    count_include_pad: bool = True,
     layout: str = "NCW",
     out_layout: Optional[str] = None,
 ) -> Expr:
@@ -1008,7 +1008,7 @@ def avg_pool3d(
     padding: Union[int, Tuple[int, ...]] = (0, 0, 0),
     dilation: Union[int, Tuple[int, int]] = (1, 1, 1),
     ceil_mode: bool = False,
-    count_include_pad: bool = False,
+    count_include_pad: bool = True,
     layout: str = "NCDHW",
     out_layout: Optional[str] = None,
 ) -> Expr: