Skip to content

[FlagGems Operator Development Competition] Add leaky_relu, leaky_relu_, and leaky_relu.out #2656

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Karl-XZ wants to merge 1 commit into from
Closed

[FlagGems Operator Development Competition] Add leaky_relu, leaky_relu_, and leaky_relu.out #2656

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
208 changes: 127 additions & 81 deletions src/flag_gems/ops/leaky_relu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,103 +5,149 @@
import triton.language as tl

from flag_gems.runtime import torch_device_fn
from flag_gems.utils import libentry
from flag_gems.utils import pointwise_dynamic

logger = logging.getLogger(__name__)

_FALLBACK_KEYSET = torch._C.DispatchKeySet(
torch._C.DispatchKey.CompositeExplicitAutograd
)

def _leaky_relu_autotune_configs():
return [
# Tiny tensors (n <= 32K): small blocks
triton.Config({"BLOCK_SIZE": 256}, num_warps=4, num_stages=2),
triton.Config({"BLOCK_SIZE": 256}, num_warps=8, num_stages=2),
triton.Config({"BLOCK_SIZE": 512}, num_warps=4, num_stages=2),
triton.Config({"BLOCK_SIZE": 512}, num_warps=8, num_stages=2),
# Small-medium tensors (n ~ 64K-4M): 1024-element blocks
triton.Config({"BLOCK_SIZE": 1024}, num_warps=4, num_stages=2),
triton.Config({"BLOCK_SIZE": 1024}, num_warps=8, num_stages=2),
triton.Config({"BLOCK_SIZE": 1024}, num_warps=4, num_stages=3),
triton.Config({"BLOCK_SIZE": 1024}, num_warps=8, num_stages=3),
triton.Config({"BLOCK_SIZE": 1024}, num_warps=4, num_stages=4),
triton.Config({"BLOCK_SIZE": 1024}, num_warps=8, num_stages=4),
triton.Config({"BLOCK_SIZE": 1024}, num_warps=16, num_stages=4),
# Medium-large tensors (n ~ 4M-16M): 2048-element blocks
triton.Config({"BLOCK_SIZE": 2048}, num_warps=4, num_stages=3),
triton.Config({"BLOCK_SIZE": 2048}, num_warps=8, num_stages=3),
triton.Config({"BLOCK_SIZE": 2048}, num_warps=8, num_stages=4),
triton.Config({"BLOCK_SIZE": 2048}, num_warps=16, num_stages=4),
triton.Config({"BLOCK_SIZE": 2048}, num_warps=4, num_stages=5),
triton.Config({"BLOCK_SIZE": 2048}, num_warps=8, num_stages=5),
triton.Config({"BLOCK_SIZE": 2048}, num_warps=16, num_stages=5),
# Large tensors (n >= 16M): 4096-element blocks for max bandwidth
triton.Config({"BLOCK_SIZE": 4096}, num_warps=4, num_stages=3),
triton.Config({"BLOCK_SIZE": 4096}, num_warps=8, num_stages=3),
triton.Config({"BLOCK_SIZE": 4096}, num_warps=8, num_stages=4),
triton.Config({"BLOCK_SIZE": 4096}, num_warps=16, num_stages=4),
triton.Config({"BLOCK_SIZE": 4096}, num_warps=4, num_stages=5),
triton.Config({"BLOCK_SIZE": 4096}, num_warps=8, num_stages=5),
triton.Config({"BLOCK_SIZE": 4096}, num_warps=16, num_stages=5),
]


@libentry()
@triton.autotune(configs=_leaky_relu_autotune_configs(), key=["n_elements"])
@triton.jit(do_not_specialize=["negative_slope"])
def _leaky_relu_kernel(
input_ptr,
output_ptr,
n_elements,
negative_slope,
BLOCK_SIZE: tl.constexpr,

@pointwise_dynamic(is_tensor=[True, False], promotion_methods=[(0, "DEFAULT")])
@triton.jit
def leaky_relu_fallback(x, negative_slope):
return tl.where(x > 0, x, x * negative_slope)


@triton.autotune(
configs=[
triton.Config({"BLOCK_SIZE": 1024}, num_warps=4),
triton.Config({"BLOCK_SIZE": 1024}, num_warps=8),
triton.Config({"BLOCK_SIZE": 2048}, num_warps=8),
triton.Config({"BLOCK_SIZE": 4096}, num_warps=4),
triton.Config({"BLOCK_SIZE": 4096}, num_warps=8),
triton.Config({"BLOCK_SIZE": 8192}, num_warps=4),
triton.Config({"BLOCK_SIZE": 8192}, num_warps=8),
triton.Config({"BLOCK_SIZE": 16384}, num_warps=4),
triton.Config({"BLOCK_SIZE": 16384}, num_warps=8),
],
key=["n_elements"],
)
@triton.jit
def leaky_relu_kernel_fp16(
x_ptr, out_ptr, negative_slope, n_elements, BLOCK_SIZE: tl.constexpr
):
pid = tl.program_id(0)
offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
pid = tl.program_id(axis=0)
block_start = pid * BLOCK_SIZE
offsets = block_start + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_elements

x = tl.load(input_ptr + offsets, mask=mask)
output = tl.where(x >= 0, x, x * negative_slope)
tl.store(output_ptr + offsets, output, mask=mask)
x = tl.load(x_ptr + offsets, mask=mask)
y = tl.where(x > 0, x, x * negative_slope)
tl.store(out_ptr + offsets, y, mask=mask)


def leaky_relu(A, negative_slope=0.01):
logger.debug("GEMS LEAKY_RELU")
if not A.is_contiguous():
A = A.contiguous()
output = torch.empty_like(A)
n_elements = A.numel()
@triton.autotune(
configs=[
triton.Config({"BLOCK_SIZE": 1024}, num_warps=4),
triton.Config({"BLOCK_SIZE": 1024}, num_warps=8),
triton.Config({"BLOCK_SIZE": 4096}, num_warps=8),
triton.Config({"BLOCK_SIZE": 8192}, num_warps=8),
triton.Config({"BLOCK_SIZE": 16384}, num_warps=4),
triton.Config({"BLOCK_SIZE": 32768}, num_warps=8),
],
key=["n_elements"],
)
@triton.jit
def leaky_relu_kernel_fp32(
x_ptr, out_ptr, negative_slope, n_elements, BLOCK_SIZE: tl.constexpr
):
pid = tl.program_id(axis=0)
block_start = pid * BLOCK_SIZE
offsets = block_start + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_elements

x = tl.load(x_ptr + offsets, mask=mask)
y = tl.where(x > 0, x, x * negative_slope)
tl.store(out_ptr + offsets, y, mask=mask)


def _get_fast_kernel(inp):
if inp.dtype in (torch.float16, torch.bfloat16):
return leaky_relu_kernel_fp16
if inp.dtype == torch.float32:
return leaky_relu_kernel_fp32
return None


def _leaky_relu_contiguous(inp, negative_slope, out):
n_elements = inp.numel()
if n_elements == 0:
return output
return out
kernel = _get_fast_kernel(inp)
if kernel is None:
return leaky_relu_fallback(inp, negative_slope, out0=out)

grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
with torch_device_fn.device(A.device.index):
_leaky_relu_kernel[grid](A, output, n_elements, negative_slope)
return output
with torch_device_fn.device(inp.device.index):
kernel[grid](
inp,
out,
negative_slope,
n_elements,
)
return out


def leaky_relu_(A, negative_slope=0.01):
def _can_use_fast_path(inp):
return (
inp.layout == torch.strided
and inp.is_cuda
and not inp.is_quantized
and not inp.is_complex()
and inp.is_contiguous()
)


def leaky_relu(inp, negative_slope=0.01):
logger.debug("GEMS LEAKY_RELU")
if _can_use_fast_path(inp):
return _leaky_relu_contiguous(inp, negative_slope, torch.empty_like(inp))
if not inp.is_cuda or inp.is_complex():
return torch.ops.aten.leaky_relu.default.redispatch(
_FALLBACK_KEYSET, inp, negative_slope
)
return leaky_relu_fallback(inp, negative_slope)


def leaky_relu_(inp, negative_slope=0.01):
logger.debug("GEMS LEAKY_RELU_")
if not A.is_contiguous():
raise RuntimeError(
"leaky_relu_ requires a contiguous tensor for in-place operation"
if _can_use_fast_path(inp):
return _leaky_relu_contiguous(inp, negative_slope, inp)
if not inp.is_cuda or inp.is_complex():
return torch.ops.aten.leaky_relu_.default.redispatch(
_FALLBACK_KEYSET, inp, negative_slope
)
n_elements = A.numel()
if n_elements == 0:
return A
grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
with torch_device_fn.device(A.device.index):
_leaky_relu_kernel[grid](A, A, n_elements, negative_slope)
return A
return leaky_relu_fallback(inp, negative_slope, out0=inp)


def leaky_relu_out(A, negative_slope=0.01, *, out=None):
def leaky_relu_out(inp, negative_slope=0.01, *, out):
logger.debug("GEMS LEAKY_RELU_OUT")
if out is None:
return leaky_relu(A, negative_slope)
if not A.is_contiguous():
A = A.contiguous()
n_elements = A.numel()
if n_elements == 0:
return out
grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
with torch_device_fn.device(A.device.index):
_leaky_relu_kernel[grid](A, out, n_elements, negative_slope)
if (
not _can_use_fast_path(inp)
or out.layout != torch.strided
or out.device != inp.device
or out.dtype != inp.dtype
):
return torch.ops.aten.leaky_relu.out.redispatch(
_FALLBACK_KEYSET, inp, negative_slope, out=out
)

if out.shape != inp.shape:
out.resize_(inp.shape)

if out.is_contiguous():
return _leaky_relu_contiguous(inp, negative_slope, out)
leaky_relu_fallback(inp, negative_slope, out0=out)
return out
45 changes: 45 additions & 0 deletions tests/test_unary_pointwise_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import warnings

import pytest
import torch

Expand Down Expand Up @@ -767,6 +769,49 @@ def test_accuracy_relu_(shape, dtype):
gems_assert_close(res_out, ref_out, dtype)


@pytest.mark.leaky_relu_out
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
def test_leaky_relu_out_resizes_output(dtype):
inp = torch.randn((8, 16), dtype=dtype, device=flag_gems.device)
ref_inp = to_reference(inp, True)
out = torch.empty((32,), dtype=dtype, device=flag_gems.device)
ref_out = torch.empty((32,), dtype=ref_inp.dtype, device=ref_inp.device)

with warnings.catch_warnings():
warnings.simplefilter("ignore", UserWarning)
torch.ops.aten.leaky_relu.out(ref_inp, 0.01, out=ref_out)
with flag_gems.use_gems():
torch.ops.aten.leaky_relu.out(inp, 0.01, out=out)

assert out.shape == ref_out.shape
gems_assert_close(out, ref_out, dtype)


@pytest.mark.leaky_relu_out
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
def test_leaky_relu_out_dtype_mismatch(dtype):
inp = torch.randn((8, 16), dtype=dtype, device=flag_gems.device)
out_dtype = torch.float16 if dtype != torch.float16 else torch.float32
out = torch.empty_like(inp, dtype=out_dtype)

with flag_gems.use_gems(), pytest.raises(RuntimeError):
torch.ops.aten.leaky_relu.out(inp, 0.01, out=out)


@pytest.mark.relu6
@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
def test_relu6(shape, dtype):
res_inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
ref_inp = to_reference(res_inp, True)

ref_out = torch.nn.functional.relu6(ref_inp)
with flag_gems.use_gems():
res_out = torch.nn.functional.relu6(res_inp)

gems_assert_close(res_out, ref_out, dtype)


@pytest.mark.softplus
@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
Expand Down