forked from ccf-ai-infra/TileOPs-Metax
496 lines
15 KiB
Python
496 lines
15 KiB
Python
import functools
|
|
import itertools
|
|
from typing import Optional, Tuple
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import torch
|
|
|
|
from tileops.kernels.kernel_base import Kernel
|
|
from tileops.kernels.pool.common import pool_output_dim
|
|
|
|
__all__ = ["MaxPool2dKernel", "MaxPool2dWithIndicesKernel"]
|
|
|
|
|
|
@functools.lru_cache(maxsize=32)
|
|
def _max_pool2d_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,
|
|
dilation_h: int,
|
|
dilation_w: int,
|
|
ceil_mode: bool,
|
|
dtype: str = "float16",
|
|
):
|
|
accum_dtype = "float"
|
|
out_h = pool_output_dim(h_in, kernel_h, stride_h, pad_h, ceil_mode, dilation_h)
|
|
out_w = pool_output_dim(w_in, kernel_w, stride_w, pad_w, ceil_mode, dilation_w)
|
|
total_output = n * c_in * out_h * out_w
|
|
|
|
@tilelang.jit(out_idx=[1], compile_flags=["-O3", "-DENABLE_BF16"])
|
|
def _max_pool2d_func(block_m: int, threads: int):
|
|
@T.prim_func
|
|
def _max_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(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
|
|
|
|
max_val = T.alloc_var(T.float32)
|
|
has_nan = T.alloc_var(T.bool)
|
|
max_val = T.cast(float("-inf"), accum_dtype)
|
|
has_nan = False
|
|
for kh in T.serial(kernel_h):
|
|
for kw in T.serial(kernel_w):
|
|
ih = oh * stride_h - pad_h + kh * dilation_h
|
|
iw = ow * stride_w - pad_w + kw * dilation_w
|
|
if ih >= 0 and ih < h_in and iw >= 0 and iw < w_in:
|
|
val = T.cast(x[batch, c_idx, ih, iw], accum_dtype)
|
|
has_nan = has_nan | T.isnan(val)
|
|
max_val = T.max(max_val, val)
|
|
|
|
result = T.if_then_else(
|
|
has_nan,
|
|
T.cast(float("nan"), accum_dtype),
|
|
max_val,
|
|
)
|
|
out[batch, c_idx, oh, ow] = T.cast(result, dtype)
|
|
|
|
return _max_pool2d_main
|
|
|
|
return _max_pool2d_func
|
|
|
|
|
|
@torch.library.custom_op("top::max_pool2d_wrapped_kernel", mutates_args=())
|
|
def _max_pool2d_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,
|
|
dilation_h: int,
|
|
dilation_w: int,
|
|
ceil_mode: bool,
|
|
dtype: str,
|
|
block_m: int,
|
|
threads: int,
|
|
x: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
return _max_pool2d_kernel(
|
|
n,
|
|
c_in,
|
|
h_in,
|
|
w_in,
|
|
kernel_h,
|
|
kernel_w,
|
|
stride_h,
|
|
stride_w,
|
|
pad_h,
|
|
pad_w,
|
|
dilation_h,
|
|
dilation_w,
|
|
ceil_mode,
|
|
dtype,
|
|
)(block_m, threads)(x)
|
|
|
|
|
|
@_max_pool2d_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,
|
|
dilation_h: int,
|
|
dilation_w: int,
|
|
ceil_mode: bool,
|
|
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, ceil_mode, dilation_h)
|
|
out_w = pool_output_dim(w_in, kernel_w, stride_w, pad_w, ceil_mode, dilation_w)
|
|
return torch.empty((n, c_in, out_h, out_w), dtype=x.dtype, device=x.device)
|
|
|
|
|
|
class MaxPool2dKernel(Kernel):
|
|
"""Max pooling forward kernel for NCHW inputs (return_indices=False)."""
|
|
|
|
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,
|
|
dilation_h: int,
|
|
dilation_w: int,
|
|
ceil_mode: bool,
|
|
dtype: torch.dtype,
|
|
config: Optional[dict] = None,
|
|
tune: bool = False,
|
|
) -> None:
|
|
super().__init__()
|
|
if dtype not in {torch.float16, torch.bfloat16, torch.float32}:
|
|
raise ValueError(
|
|
f"MaxPool2dKernel supports float16, bfloat16, and float32, got {dtype}"
|
|
)
|
|
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.dilation_h = dilation_h
|
|
self.dilation_w = dilation_w
|
|
self.ceil_mode = ceil_mode
|
|
self.dtype = dtype
|
|
self.out_h = pool_output_dim(h_in, kernel_h, stride_h, pad_h, ceil_mode, dilation_h)
|
|
self.out_w = pool_output_dim(w_in, kernel_w, stride_w, pad_w, ceil_mode, dilation_w)
|
|
self.kernel = _max_pool2d_kernel(
|
|
n,
|
|
c_in,
|
|
h_in,
|
|
w_in,
|
|
kernel_h,
|
|
kernel_w,
|
|
stride_h,
|
|
stride_w,
|
|
pad_h,
|
|
pad_w,
|
|
dilation_h,
|
|
dilation_w,
|
|
ceil_mode,
|
|
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], [128, 256, 512])
|
|
]
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
return _max_pool2d_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.dilation_h,
|
|
self.dilation_w,
|
|
self.ceil_mode,
|
|
self.dtype_str,
|
|
self.config["block_m"],
|
|
self.config["threads"],
|
|
x,
|
|
)
|
|
|
|
|
|
@functools.lru_cache(maxsize=32)
|
|
def _max_pool2d_with_indices_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,
|
|
dilation_h: int,
|
|
dilation_w: int,
|
|
ceil_mode: bool,
|
|
dtype: str = "float16",
|
|
):
|
|
accum_dtype = "float"
|
|
out_h = pool_output_dim(h_in, kernel_h, stride_h, pad_h, ceil_mode, dilation_h)
|
|
out_w = pool_output_dim(w_in, kernel_w, stride_w, pad_w, ceil_mode, dilation_w)
|
|
total_output = n * c_in * out_h * out_w
|
|
|
|
@tilelang.jit(out_idx=[1, 2], compile_flags=["-O3", "-DENABLE_BF16"])
|
|
def _max_pool2d_with_indices_func(block_m: int, threads: int):
|
|
@T.prim_func
|
|
def _max_pool2d_with_indices_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
|
|
indices: T.Tensor((n, c_in, out_h, out_w), "int64"), # 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
|
|
|
|
max_val = T.alloc_var(T.float32)
|
|
has_nan = T.alloc_var(T.bool)
|
|
max_idx = T.alloc_var(T.int64)
|
|
nan_idx = T.alloc_var(T.int64)
|
|
first_valid = T.alloc_var(T.bool)
|
|
max_val = T.cast(float("-inf"), accum_dtype)
|
|
has_nan = False
|
|
max_idx = T.cast(0, "int64")
|
|
nan_idx = T.cast(0, "int64")
|
|
first_valid = True
|
|
for kh in T.serial(kernel_h):
|
|
for kw in T.serial(kernel_w):
|
|
ih = oh * stride_h - pad_h + kh * dilation_h
|
|
iw = ow * stride_w - pad_w + kw * dilation_w
|
|
if ih >= 0 and ih < h_in and iw >= 0 and iw < w_in:
|
|
val = T.cast(x[batch, c_idx, ih, iw], accum_dtype)
|
|
flat_idx = T.cast(ih * w_in + iw, "int64")
|
|
is_nan = T.isnan(val)
|
|
if is_nan:
|
|
# PyTorch records the last NaN visited in
|
|
# a pooling window.
|
|
nan_idx = flat_idx
|
|
has_nan = True
|
|
elif first_valid:
|
|
max_val = val
|
|
max_idx = flat_idx
|
|
first_valid = False
|
|
elif val > max_val:
|
|
max_val = val
|
|
max_idx = flat_idx
|
|
|
|
result = T.if_then_else(
|
|
has_nan,
|
|
T.cast(float("nan"), accum_dtype),
|
|
max_val,
|
|
)
|
|
out[batch, c_idx, oh, ow] = T.cast(result, dtype)
|
|
indices[batch, c_idx, oh, ow] = T.if_then_else(
|
|
has_nan,
|
|
nan_idx,
|
|
max_idx,
|
|
)
|
|
|
|
return _max_pool2d_with_indices_main
|
|
|
|
return _max_pool2d_with_indices_func
|
|
|
|
|
|
@torch.library.custom_op("top::max_pool2d_with_indices_wrapped_kernel", mutates_args=())
|
|
def _max_pool2d_with_indices_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,
|
|
dilation_h: int,
|
|
dilation_w: int,
|
|
ceil_mode: bool,
|
|
dtype: str,
|
|
block_m: int,
|
|
threads: int,
|
|
x: torch.Tensor,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
return _max_pool2d_with_indices_kernel(
|
|
n,
|
|
c_in,
|
|
h_in,
|
|
w_in,
|
|
kernel_h,
|
|
kernel_w,
|
|
stride_h,
|
|
stride_w,
|
|
pad_h,
|
|
pad_w,
|
|
dilation_h,
|
|
dilation_w,
|
|
ceil_mode,
|
|
dtype,
|
|
)(block_m, threads)(x)
|
|
|
|
|
|
@_max_pool2d_with_indices_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,
|
|
dilation_h: int,
|
|
dilation_w: int,
|
|
ceil_mode: bool,
|
|
dtype: str,
|
|
block_m: int,
|
|
threads: int,
|
|
x: torch.Tensor,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
_ = (dtype, block_m, threads)
|
|
out_h = pool_output_dim(h_in, kernel_h, stride_h, pad_h, ceil_mode, dilation_h)
|
|
out_w = pool_output_dim(w_in, kernel_w, stride_w, pad_w, ceil_mode, dilation_w)
|
|
return (
|
|
torch.empty((n, c_in, out_h, out_w), dtype=x.dtype, device=x.device),
|
|
torch.empty((n, c_in, out_h, out_w), dtype=torch.int64, device=x.device),
|
|
)
|
|
|
|
|
|
class MaxPool2dWithIndicesKernel(Kernel):
|
|
"""Max pooling forward-with-indices kernel for NCHW inputs."""
|
|
|
|
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,
|
|
dilation_h: int,
|
|
dilation_w: int,
|
|
ceil_mode: bool,
|
|
dtype: torch.dtype,
|
|
config: Optional[dict] = None,
|
|
tune: bool = False,
|
|
) -> None:
|
|
super().__init__()
|
|
if dtype not in {torch.float16, torch.bfloat16, torch.float32}:
|
|
raise ValueError(
|
|
f"MaxPool2dWithIndicesKernel supports float16, bfloat16, and float32, got {dtype}"
|
|
)
|
|
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.dilation_h = dilation_h
|
|
self.dilation_w = dilation_w
|
|
self.ceil_mode = ceil_mode
|
|
self.dtype = dtype
|
|
self.out_h = pool_output_dim(h_in, kernel_h, stride_h, pad_h, ceil_mode, dilation_h)
|
|
self.out_w = pool_output_dim(w_in, kernel_w, stride_w, pad_w, ceil_mode, dilation_w)
|
|
self.kernel = _max_pool2d_with_indices_kernel(
|
|
n,
|
|
c_in,
|
|
h_in,
|
|
w_in,
|
|
kernel_h,
|
|
kernel_w,
|
|
stride_h,
|
|
stride_w,
|
|
pad_h,
|
|
pad_w,
|
|
dilation_h,
|
|
dilation_w,
|
|
ceil_mode,
|
|
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], [128, 256, 512])
|
|
]
|
|
|
|
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
return _max_pool2d_with_indices_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.dilation_h,
|
|
self.dilation_w,
|
|
self.ceil_mode,
|
|
self.dtype_str,
|
|
self.config["block_m"],
|
|
self.config["threads"],
|
|
x,
|
|
)
|