diff --git a/tests/ops/test_pool.py b/tests/ops/test_pool.py index f25abd6d1..022190d9f 100644 --- a/tests/ops/test_pool.py +++ b/tests/ops/test_pool.py @@ -6,7 +6,7 @@ from tests.test_base import FixtureBase, TestBase from tileops.kernels.kernel_base import Kernel -from tileops.kernels.pool import AvgPool2dKernel, AvgPool3dKernel +from tileops.kernels.pool import AvgPool2dSpatialKernel, AvgPool3dKernel from tileops.ops import AvgPool1dFwdOp, AvgPool2dFwdOp, AvgPool3dFwdOp @@ -186,6 +186,31 @@ class AvgPool2dFixture(FixtureBase): marks=pytest.mark.full, id="full-divisor-override-bf16", ), + pytest.param( + 1, 7, 9, 10, (3, 3), (2, 2), (1, 1), False, False, None, torch.float16, False, + marks=pytest.mark.full, + id="full-no-ceil-no-pad-count-fp16", + ), + pytest.param( + 1, 5, 10, 11, (3, 3), (2, 2), (1, 1), True, True, None, torch.float32, False, + marks=pytest.mark.full, + id="full-ceil-pad-count-fp32", + ), + pytest.param( + 2, 6, 9, 13, (2, 3), (2, 2), (0, 1), True, True, 7, torch.bfloat16, False, + marks=pytest.mark.full, + id="full-ceil-pad-count-divisor-bf16", + ), + pytest.param( + 1, 9, 11, 12, (3, 5), (2, 3), (1, 2), True, False, 7, torch.float16, False, + marks=pytest.mark.full, + id="full-ceil-no-pad-count-divisor-fp16", + ), + pytest.param( + 1, 8, 9, 9, (3, 3), (2, 2), (1, 1), False, False, 7, torch.float16, False, + marks=pytest.mark.full, + id="full-no-ceil-no-pad-count-divisor-fp16", + ), ]), ] @@ -279,7 +304,23 @@ def test_avg_pool2d_dispatches_kernel() -> None: stride=(2, 2), padding=(1, 1), ) - assert isinstance(op.kernel, AvgPool2dKernel) + assert isinstance(op.kernel, AvgPool2dSpatialKernel) + + +@pytest.mark.smoke +def test_avg_pool2d_rejects_non_positive_output_size() -> None: + with pytest.raises(ValueError, match="output size must be greater than zero"): + AvgPool2dFwdOp( + n=1, + c_in=1, + h_in=2, + w_in=2, + kernel_size=(5, 5), + stride=(1, 1), + padding=(0, 0), + ceil_mode=False, + count_include_pad=True, + ) @pytest.mark.smoke diff --git a/tileops/kernels/pool/__init__.py b/tileops/kernels/pool/__init__.py index bf026df3d..40526a396 100644 --- a/tileops/kernels/pool/__init__.py +++ b/tileops/kernels/pool/__init__.py @@ -1,5 +1,10 @@ from .avg_pool1d import AvgPool1dKernel -from .avg_pool2d import AvgPool2dKernel +from .avg_pool2d import AvgPool2dKernel, AvgPool2dSpatialKernel from .avg_pool3d import AvgPool3dKernel -__all__ = ["AvgPool1dKernel", "AvgPool2dKernel", "AvgPool3dKernel"] +__all__ = [ + "AvgPool1dKernel", + "AvgPool2dKernel", + "AvgPool2dSpatialKernel", + "AvgPool3dKernel", +] diff --git a/tileops/kernels/pool/avg_pool2d.py b/tileops/kernels/pool/avg_pool2d.py index e2315d072..d15495b94 100644 --- a/tileops/kernels/pool/avg_pool2d.py +++ b/tileops/kernels/pool/avg_pool2d.py @@ -9,7 +9,7 @@ from tileops.kernels.kernel_base import Kernel from tileops.kernels.pool.common import pool_output_dim -__all__ = ["AvgPool2dKernel"] +__all__ = ["AvgPool2dKernel", "AvgPool2dSpatialKernel"] @functools.lru_cache(maxsize=64) @@ -33,108 +33,189 @@ def _avg_pool2d_kernel( accum_dtype = "float" out_h = pool_output_dim(h_in, kernel_h, stride_h, pad_h, ceil_mode) out_w = pool_output_dim(w_in, kernel_w, stride_w, pad_w, ceil_mode) + total_output = n * c_in * out_h * out_w @tilelang.jit(out_idx=[1], compile_flags=["-O3", "-DENABLE_BF16"]) - def _avg_pool2d_func(block_m: int, block_c: int, threads: int): + def _avg_pool2d_func(block_m: int, threads: int): @T.prim_func def _avg_pool2d_main( x: T.Tensor((n, c_in, h_in, w_in), dtype), # type: ignore out: T.Tensor((n, c_in, out_h, out_w), dtype), # type: ignore ): - with T.Kernel( - T.ceildiv(out_h * out_w, block_m), - T.ceildiv(c_in, block_c), - n, - threads=threads, - ) as (bx, by, bz): - T.use_swizzle(10) - tile_spatial_start = bx * block_m - tile_spatial_end = tile_spatial_start + block_m - 1 - tile_oh_start = tile_spatial_start // out_w - tile_oh_end = tile_spatial_end // out_w - tile_spatial_same_row = tile_oh_start == tile_oh_end - tile_ow_start = tile_spatial_start % out_w - tile_ow_end = tile_spatial_end % out_w - tile_input_h_start = tile_oh_start * stride_h - pad_h - tile_input_h_end = tile_oh_end * stride_h + kernel_h - 1 - pad_h - tile_input_w_start = tile_ow_start * stride_w - pad_w - tile_input_w_end = tile_ow_end * stride_w + kernel_w - 1 - pad_w - tile_spatial_full = ( - tile_spatial_same_row - & (tile_input_h_start >= 0) - & (tile_input_h_end < h_in) - & (tile_input_w_start >= 0) - & (tile_input_w_end < w_in) - ) - use_fixed_kernel_divisor = count_include_pad and not use_divisor_override - - for i, j in T.Parallel(block_m, block_c): - out_spatial = bx * block_m + i - oh = out_spatial // out_w - ow = out_spatial % out_w - c_idx = by * block_c + j - batch = bz - if out_spatial < out_h * out_w and c_idx < c_in: + with T.Kernel(T.ceildiv(total_output, block_m), threads=threads) as bx: + for i in T.Parallel(block_m): + out_idx = bx * block_m + i + if out_idx < total_output: + ow = out_idx % out_w + spatial_idx = out_idx // out_w + oh = spatial_idx % out_h + channel_batch_idx = spatial_idx // out_h + c_idx = channel_batch_idx % c_in + batch = channel_batch_idx // c_in + sum_val = T.alloc_var(T.float32) sum_val = T.cast(0.0, accum_dtype) - - if tile_spatial_full and use_fixed_kernel_divisor: - for kh in T.serial(kernel_h): - for kw in T.serial(kernel_w): - ih = oh * stride_h + kh - pad_h - iw = ow * stride_w + kw - pad_w + for kh in T.serial(kernel_h): + for kw in T.serial(kernel_w): + ih = oh * stride_h + kh - pad_h + iw = ow * stride_w + kw - pad_w + if ih >= 0 and ih < h_in and iw >= 0 and iw < w_in: sum_val += T.cast( x[batch, c_idx, ih, iw], accum_dtype ) - out[batch, c_idx, oh, ow] = T.cast( - sum_val / T.cast(kernel_h * kernel_w, accum_dtype), - dtype, - ) - else: - for kh in T.serial(kernel_h): - for kw in T.serial(kernel_w): - ih = oh * stride_h + kh - pad_h - iw = ow * stride_w + kw - pad_w - if ih >= 0 and ih < h_in and iw >= 0 and iw < w_in: - sum_val += T.cast( - x[batch, c_idx, ih, iw], accum_dtype - ) - - start_h = oh * stride_h - pad_h - start_w = ow * stride_w - pad_w - end_h = start_h + kernel_h - end_w = start_w + kernel_w - valid_h = T.max(T.min(end_h, h_in) - T.max(start_h, 0), 0) - valid_w = T.max(T.min(end_w, w_in) - T.max(start_w, 0), 0) - valid_count = valid_h * valid_w - padded_h = T.max( - T.min(end_h, h_in + pad_h) - T.max(start_h, -pad_h), 0 - ) - padded_w = T.max( - T.min(end_w, w_in + pad_w) - T.max(start_w, -pad_w), 0 - ) - padded_count = padded_h * padded_w - auto_divisor = T.max( - T.if_then_else( - count_include_pad, padded_count, valid_count - ), - 1, - ) - divisor = T.if_then_else( - use_divisor_override, - divisor_override, - auto_divisor, - ) - out[batch, c_idx, oh, ow] = T.cast( - sum_val / T.cast(divisor, accum_dtype), - dtype, - ) + + start_h = oh * stride_h - pad_h + start_w = ow * stride_w - pad_w + end_h = start_h + kernel_h + end_w = start_w + kernel_w + valid_h = T.max(T.min(end_h, h_in) - T.max(start_h, 0), 0) + valid_w = T.max(T.min(end_w, w_in) - T.max(start_w, 0), 0) + valid_count = valid_h * valid_w + padded_h = T.max( + T.min(end_h, h_in + pad_h) - T.max(start_h, -pad_h), 0 + ) + padded_w = T.max( + T.min(end_w, w_in + pad_w) - T.max(start_w, -pad_w), 0 + ) + padded_count = padded_h * padded_w + auto_divisor = T.max( + T.if_then_else( + count_include_pad, padded_count, valid_count + ), + 1, + ) + divisor = T.if_then_else( + use_divisor_override, + divisor_override, + auto_divisor, + ) + out[batch, c_idx, oh, ow] = T.cast( + sum_val / T.cast(divisor, accum_dtype), + dtype, + ) return _avg_pool2d_main return _avg_pool2d_func +@functools.lru_cache(maxsize=64) +def _avg_pool2d_spatial_kernel( + n: int, + c_in: int, + h_in: int, + w_in: int, + kernel_h: int, + kernel_w: int, + stride_h: int, + stride_w: int, + pad_h: int, + pad_w: int, + dtype: str = "float16", +): + accum_dtype = "float" + out_h = pool_output_dim(h_in, kernel_h, stride_h, pad_h, False) + out_w = pool_output_dim(w_in, kernel_w, stride_w, pad_w, False) + total_output = n * c_in * out_h * out_w + + @tilelang.jit(out_idx=[1], compile_flags=["-O3", "-DENABLE_BF16"]) + def _avg_pool2d_spatial_func(block_m: int, threads: int): + @T.prim_func + def _avg_pool2d_spatial_main( + x: T.Tensor((n, c_in, h_in, w_in), dtype), # type: ignore + out: T.Tensor((n, c_in, out_h, out_w), dtype), # type: ignore + ): + with T.Kernel(T.ceildiv(total_output, block_m), threads=threads) as bx: + for i in T.Parallel(block_m): + out_idx = bx * block_m + i + if out_idx < total_output: + ow = out_idx % out_w + spatial_idx = out_idx // out_w + oh = spatial_idx % out_h + channel_batch_idx = spatial_idx // out_h + c_idx = channel_batch_idx % c_in + batch = channel_batch_idx // c_in + + sum_val = T.alloc_var(T.float32) + sum_val = T.cast(0.0, accum_dtype) + for kh in T.serial(kernel_h): + for kw in T.serial(kernel_w): + ih = oh * stride_h + kh - pad_h + iw = ow * stride_w + kw - pad_w + if ih >= 0 and ih < h_in and iw >= 0 and iw < w_in: + sum_val += T.cast( + x[batch, c_idx, ih, iw], accum_dtype + ) + + out[batch, c_idx, oh, ow] = T.cast( + sum_val / T.cast(kernel_h * kernel_w, accum_dtype), + dtype, + ) + + return _avg_pool2d_spatial_main + + return _avg_pool2d_spatial_func + + +@torch.library.custom_op("top::avg_pool2d_spatial_wrapped_kernel", mutates_args=()) +def _avg_pool2d_spatial_wrapped_kernel( + n: int, + c_in: int, + h_in: int, + w_in: int, + kernel_h: int, + kernel_w: int, + stride_h: int, + stride_w: int, + pad_h: int, + pad_w: int, + dtype: str, + block_m: int, + threads: int, + x: torch.Tensor, +) -> torch.Tensor: + return _avg_pool2d_spatial_kernel( + n, + c_in, + h_in, + w_in, + kernel_h, + kernel_w, + stride_h, + stride_w, + pad_h, + pad_w, + dtype, + )(block_m, threads)(x) + + +@_avg_pool2d_spatial_wrapped_kernel.register_fake +def _( + n: int, + c_in: int, + h_in: int, + w_in: int, + kernel_h: int, + kernel_w: int, + stride_h: int, + stride_w: int, + pad_h: int, + pad_w: int, + dtype: str, + block_m: int, + threads: int, + x: torch.Tensor, +) -> torch.Tensor: + _ = ( + dtype, + block_m, + threads, + ) + out_h = pool_output_dim(h_in, kernel_h, stride_h, pad_h, False) + out_w = pool_output_dim(w_in, kernel_w, stride_w, pad_w, False) + return torch.empty((n, c_in, out_h, out_w), dtype=x.dtype, device=x.device) + + @torch.library.custom_op("top::avg_pool2d_wrapped_kernel", mutates_args=()) def _avg_pool2d_wrapped_kernel( n: int, @@ -153,7 +234,6 @@ def _avg_pool2d_wrapped_kernel( divisor_override: int, dtype: str, block_m: int, - block_c: int, threads: int, x: torch.Tensor, ) -> torch.Tensor: @@ -173,7 +253,7 @@ def _avg_pool2d_wrapped_kernel( use_divisor_override, divisor_override, dtype, - )(block_m, block_c, threads)(x) + )(block_m, threads)(x) @_avg_pool2d_wrapped_kernel.register_fake @@ -194,7 +274,6 @@ def _( divisor_override: int, dtype: str, block_m: int, - block_c: int, threads: int, x: torch.Tensor, ) -> torch.Tensor: @@ -204,7 +283,6 @@ def _( divisor_override, dtype, block_m, - block_c, threads, ) out_h = pool_output_dim(h_in, kernel_h, stride_h, pad_h, ceil_mode) @@ -212,6 +290,89 @@ def _( return torch.empty((n, c_in, out_h, out_w), dtype=x.dtype, device=x.device) +class AvgPool2dSpatialKernel(Kernel): + """Fast path for common NCHW avg_pool2d workloads.""" + + supported_archs: list[int] = [80, 86, 89, 90] + + def __init__( + self, + n: int, + c_in: int, + h_in: int, + w_in: int, + kernel_h: int, + kernel_w: int, + stride_h: int, + stride_w: int, + pad_h: int, + pad_w: int, + dtype: torch.dtype, + config: Optional[dict] = None, + tune: bool = False, + ) -> None: + super().__init__() + self.n = n + self.c_in = c_in + self.h_in = h_in + self.w_in = w_in + self.kernel_h = kernel_h + self.kernel_w = kernel_w + self.stride_h = stride_h + self.stride_w = stride_w + self.pad_h = pad_h + self.pad_w = pad_w + self.dtype = dtype + self.out_h = pool_output_dim(h_in, kernel_h, stride_h, pad_h, False) + self.out_w = pool_output_dim(w_in, kernel_w, stride_w, pad_w, False) + self.kernel = _avg_pool2d_spatial_kernel( + n, + c_in, + h_in, + w_in, + kernel_h, + kernel_w, + stride_h, + stride_w, + pad_h, + pad_w, + self.dtype_str, + ) + self.init_config(config, tune) + + @property + def default_config(self) -> dict: + return { + "block_m": 256, + "threads": 256, + } + + @property + def autotune_configs(self) -> list[dict]: + return [ + {"block_m": block_m, "threads": threads} + for block_m, threads in itertools.product([128, 256, 512], [256, 512]) + ] + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return _avg_pool2d_spatial_wrapped_kernel( + self.n, + self.c_in, + self.h_in, + self.w_in, + self.kernel_h, + self.kernel_w, + self.stride_h, + self.stride_w, + self.pad_h, + self.pad_w, + self.dtype_str, + self.config["block_m"], + self.config["threads"], + x, + ) + + class AvgPool2dKernel(Kernel): supported_archs: list[int] = [80, 86, 89, 90] @@ -275,21 +436,15 @@ def __init__( @property def default_config(self) -> dict: return { - "block_m": 128, - "block_c": 64, - "threads": 128, + "block_m": 256, + "threads": 256, } @property def autotune_configs(self) -> list[dict]: - configs = itertools.product([64, 128, 256], [32, 64, 128], [128, 256]) return [ - { - "block_m": block_m, - "block_c": block_c, - "threads": threads, - } - for block_m, block_c, threads in configs + {"block_m": block_m, "threads": threads} + for block_m, threads in itertools.product([128, 256, 512], [128, 256, 512]) ] def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -310,7 +465,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: self.divisor_override, self.dtype_str, self.config["block_m"], - self.config["block_c"], self.config["threads"], x, ) diff --git a/tileops/ops/pool.py b/tileops/ops/pool.py index d319671a2..f3a6dcb06 100644 --- a/tileops/ops/pool.py +++ b/tileops/ops/pool.py @@ -3,7 +3,12 @@ import torch from tileops.kernels.kernel_base import Kernel -from tileops.kernels.pool import AvgPool1dKernel, AvgPool2dKernel, AvgPool3dKernel +from tileops.kernels.pool import ( + AvgPool1dKernel, + AvgPool2dKernel, + AvgPool2dSpatialKernel, + AvgPool3dKernel, +) from tileops.kernels.pool.common import ( _normalize_pool_dims, pool_output_dim, @@ -191,9 +196,13 @@ def __init__( ) self.dispatch_kernel(kernel_map) - if "avg_pool2d_kernel" not in self.kernel_map: + if ( + "avg_pool2d_kernel" not in self.kernel_map + and "avg_pool2d_spatial_kernel" not in self.kernel_map + ): raise NotImplementedError( - "AvgPool2dFwdOp requires 'avg_pool2d_kernel' in kernel_map" + "AvgPool2dFwdOp requires 'avg_pool2d_kernel' or " + "'avg_pool2d_spatial_kernel' in kernel_map" ) self.out_h = pool_output_dim( self.h_in, self.kernel_size[0], self.stride[0], self.padding[0], self.ceil_mode @@ -201,7 +210,13 @@ def __init__( self.out_w = pool_output_dim( self.w_in, self.kernel_size[1], self.stride[1], self.padding[1], self.ceil_mode ) - self.kernel = self.kernel_map["avg_pool2d_kernel"]( + if self.out_h <= 0 or self.out_w <= 0: + raise ValueError( + "AvgPool2dFwdOp calculated output size must be greater than zero, " + f"got ({self.out_h}, {self.out_w})" + ) + + kernel_kwargs = dict( n=n, c_in=c_in, h_in=h_in, @@ -212,16 +227,36 @@ def __init__( stride_w=self.stride[1], pad_h=self.padding[0], pad_w=self.padding[1], - ceil_mode=ceil_mode, - count_include_pad=count_include_pad, - divisor_override=divisor_override, dtype=dtype, tune=tune, ) + use_spatial_fast_path = ( + not ceil_mode + and count_include_pad + and divisor_override is None + and "avg_pool2d_spatial_kernel" in self.kernel_map + ) + if use_spatial_fast_path: + self.kernel = self.kernel_map["avg_pool2d_spatial_kernel"](**kernel_kwargs) + elif "avg_pool2d_kernel" in self.kernel_map: + self.kernel = self.kernel_map["avg_pool2d_kernel"]( + **kernel_kwargs, + ceil_mode=ceil_mode, + count_include_pad=count_include_pad, + divisor_override=divisor_override, + ) + else: + raise NotImplementedError( + "AvgPool2dFwdOp requires 'avg_pool2d_kernel' or " + "'avg_pool2d_spatial_kernel' in kernel_map" + ) @property def default_kernel_map(self) -> Dict[str, Kernel]: - return {"avg_pool2d_kernel": AvgPool2dKernel} + return { + "avg_pool2d_kernel": AvgPool2dKernel, + "avg_pool2d_spatial_kernel": AvgPool2dSpatialKernel, + } def _infer_output_shapes( self, input_shape: tuple[int, ...]