Skip to content

Commit d51e735

Browse files
Karl-XZKarl-XZ
committed
[FlagGems Operator Development Competition] add leaky_relu
1 parent da4466e commit d51e735

5 files changed

Lines changed: 244 additions & 0 deletions

File tree

‎benchmark/test_unary_pointwise_perf.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@ def get_input_iter(self, cur_dtype) -> Generator:
7777
("elu", torch.nn.functional.elu, FLOAT_DTYPES),
7878
("gelu", torch.nn.functional.gelu, FLOAT_DTYPES),
7979
("hardsigmoid", torch.nn.functional.hardsigmoid, FLOAT_DTYPES),
80+
("leaky_relu", torch.nn.functional.leaky_relu, FLOAT_DTYPES),
8081
("relu", torch.nn.functional.relu, FLOAT_DTYPES),
8182
("relu6", torch.nn.functional.relu6, FLOAT_DTYPES),
8283
("selu", torch.nn.functional.selu, FLOAT_DTYPES),
@@ -142,6 +143,11 @@ def test_general_unary_pointwise(op_name, torch_op, dtypes):
142143
("floor_", torch.Tensor.floor_, FLOAT_DTYPES),
143144
("gelu_", torch.ops.aten.gelu_.default, FLOAT_DTYPES),
144145
("hardswish_", torch.ops.aten.hardswish_, FLOAT_DTYPES),
146+
(
147+
"leaky_relu_",
148+
lambda a: torch.nn.functional.leaky_relu(a, inplace=True),
149+
FLOAT_DTYPES,
150+
),
145151
("log10_", torch.log10_, FLOAT_DTYPES),
146152
("neg_", torch.neg_, FLOAT_DTYPES),
147153
("reciprocal_", torch.reciprocal_, FLOAT_DTYPES),

‎src/flag_gems/__init__.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -256,6 +256,9 @@ def torch_ge(v):
256256
("kron", kron),
257257
("le.Scalar", le_scalar),
258258
("le.Tensor", le),
259+
("leaky_relu", leaky_relu),
260+
("leaky_relu_", leaky_relu_),
261+
("leaky_relu.out", leaky_relu_out),
259262
("lerp.Scalar", lerp_scalar),
260263
("lerp.Tensor", lerp_tensor),
261264
("lerp_.Scalar", lerp_scalar_),

‎src/flag_gems/ops/__init__.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,6 +160,7 @@
160160
from flag_gems.ops.kron import kron
161161
from flag_gems.ops.layernorm import layer_norm, layer_norm_backward
162162
from flag_gems.ops.le import le, le_scalar
163+
from flag_gems.ops.leaky_relu import leaky_relu, leaky_relu_, leaky_relu_out
163164
from flag_gems.ops.lerp import lerp_scalar, lerp_scalar_, lerp_tensor, lerp_tensor_
164165
from flag_gems.ops.lift_fresh_copy import lift_fresh_copy, lift_fresh_copy_out
165166
from flag_gems.ops.linspace import linspace
@@ -525,6 +526,9 @@
525526
"layer_norm_backward",
526527
"le",
527528
"le_scalar",
529+
"leaky_relu",
530+
"leaky_relu_",
531+
"leaky_relu_out",
528532
"lerp_scalar",
529533
"lerp_scalar_",
530534
"lerp_tensor",

‎src/flag_gems/ops/leaky_relu.py‎

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
import logging
2+
3+
import torch
4+
import triton
5+
import triton.language as tl
6+
7+
from flag_gems.runtime import torch_device_fn
8+
from flag_gems.utils import pointwise_dynamic
9+
10+
logger = logging.getLogger(__name__)
11+
12+
_FALLBACK_KEYSET = torch._C.DispatchKeySet(
13+
torch._C.DispatchKey.CompositeExplicitAutograd
14+
)
15+
16+
17+
@pointwise_dynamic(is_tensor=[True, False], promotion_methods=[(0, "DEFAULT")])
18+
@triton.jit
19+
def leaky_relu_fallback(x, negative_slope):
20+
return tl.where(x > 0, x, x * negative_slope)
21+
22+
23+
@triton.autotune(
24+
configs=[
25+
triton.Config({"BLOCK_SIZE": 1024}, num_warps=4),
26+
triton.Config({"BLOCK_SIZE": 1024}, num_warps=8),
27+
triton.Config({"BLOCK_SIZE": 2048}, num_warps=8),
28+
triton.Config({"BLOCK_SIZE": 4096}, num_warps=4),
29+
triton.Config({"BLOCK_SIZE": 4096}, num_warps=8),
30+
triton.Config({"BLOCK_SIZE": 8192}, num_warps=4),
31+
triton.Config({"BLOCK_SIZE": 8192}, num_warps=8),
32+
triton.Config({"BLOCK_SIZE": 16384}, num_warps=4),
33+
triton.Config({"BLOCK_SIZE": 16384}, num_warps=8),
34+
],
35+
key=["n_elements"],
36+
)
37+
@triton.jit
38+
def leaky_relu_kernel_fp16(
39+
x_ptr, out_ptr, negative_slope, n_elements, BLOCK_SIZE: tl.constexpr
40+
):
41+
pid = tl.program_id(axis=0)
42+
block_start = pid * BLOCK_SIZE
43+
offsets = block_start + tl.arange(0, BLOCK_SIZE)
44+
mask = offsets < n_elements
45+
46+
x = tl.load(x_ptr + offsets, mask=mask)
47+
y = tl.where(x > 0, x, x * negative_slope)
48+
tl.store(out_ptr + offsets, y, mask=mask)
49+
50+
51+
@triton.autotune(
52+
configs=[
53+
triton.Config({"BLOCK_SIZE": 1024}, num_warps=4),
54+
triton.Config({"BLOCK_SIZE": 1024}, num_warps=8),
55+
triton.Config({"BLOCK_SIZE": 4096}, num_warps=8),
56+
triton.Config({"BLOCK_SIZE": 8192}, num_warps=8),
57+
triton.Config({"BLOCK_SIZE": 16384}, num_warps=4),
58+
triton.Config({"BLOCK_SIZE": 32768}, num_warps=8),
59+
],
60+
key=["n_elements"],
61+
)
62+
@triton.jit
63+
def leaky_relu_kernel_fp32(
64+
x_ptr, out_ptr, negative_slope, n_elements, BLOCK_SIZE: tl.constexpr
65+
):
66+
pid = tl.program_id(axis=0)
67+
block_start = pid * BLOCK_SIZE
68+
offsets = block_start + tl.arange(0, BLOCK_SIZE)
69+
mask = offsets < n_elements
70+
71+
x = tl.load(x_ptr + offsets, mask=mask)
72+
y = tl.where(x > 0, x, x * negative_slope)
73+
tl.store(out_ptr + offsets, y, mask=mask)
74+
75+
76+
def _get_fast_kernel(inp):
77+
if inp.dtype in (torch.float16, torch.bfloat16):
78+
return leaky_relu_kernel_fp16
79+
if inp.dtype == torch.float32:
80+
return leaky_relu_kernel_fp32
81+
return None
82+
83+
84+
def _leaky_relu_contiguous(inp, negative_slope, out):
85+
n_elements = inp.numel()
86+
if n_elements == 0:
87+
return out
88+
kernel = _get_fast_kernel(inp)
89+
if kernel is None:
90+
return leaky_relu_fallback(inp, negative_slope, out0=out)
91+
92+
grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
93+
with torch_device_fn.device(inp.device):
94+
kernel[grid](
95+
inp,
96+
out,
97+
negative_slope,
98+
n_elements,
99+
)
100+
return out
101+
102+
103+
def _can_use_fast_path(inp):
104+
return (
105+
inp.layout == torch.strided
106+
and inp.is_cuda
107+
and not inp.is_quantized
108+
and not inp.is_complex()
109+
and inp.is_contiguous()
110+
)
111+
112+
113+
def leaky_relu(inp, negative_slope=0.01):
114+
logger.debug("GEMS LEAKY_RELU")
115+
if _can_use_fast_path(inp):
116+
return _leaky_relu_contiguous(inp, negative_slope, torch.empty_like(inp))
117+
if not inp.is_cuda or inp.is_complex():
118+
return torch.ops.aten.leaky_relu.default.redispatch(
119+
_FALLBACK_KEYSET, inp, negative_slope
120+
)
121+
return leaky_relu_fallback(inp, negative_slope)
122+
123+
124+
def leaky_relu_(inp, negative_slope=0.01):
125+
logger.debug("GEMS LEAKY_RELU_")
126+
if _can_use_fast_path(inp):
127+
return _leaky_relu_contiguous(inp, negative_slope, inp)
128+
if not inp.is_cuda or inp.is_complex():
129+
return torch.ops.aten.leaky_relu_.default.redispatch(
130+
_FALLBACK_KEYSET, inp, negative_slope
131+
)
132+
return leaky_relu_fallback(inp, negative_slope, out0=inp)
133+
134+
135+
def leaky_relu_out(inp, negative_slope=0.01, *, out):
136+
logger.debug("GEMS LEAKY_RELU_OUT")
137+
if (
138+
not _can_use_fast_path(inp)
139+
or out.layout != torch.strided
140+
or out.device != inp.device
141+
or out.dtype != inp.dtype
142+
):
143+
return torch.ops.aten.leaky_relu.out.redispatch(
144+
_FALLBACK_KEYSET, inp, negative_slope, out=out
145+
)
146+
147+
if out.shape != inp.shape:
148+
out.resize_(inp.shape)
149+
150+
if out.is_contiguous():
151+
return _leaky_relu_contiguous(inp, negative_slope, out)
152+
leaky_relu_fallback(inp, negative_slope, out0=out)
153+
return out

‎tests/test_unary_pointwise_ops.py‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
import warnings
2+
13
import pytest
24
import torch
35

@@ -989,6 +991,82 @@ def test_relu_(shape, dtype):
989991
gems_assert_close(res_out, ref_out, dtype)
990992

991993

994+
@pytest.mark.leaky_relu
995+
@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
996+
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
997+
@pytest.mark.parametrize("negative_slope", [0.01, 0.2, -0.5])
998+
def test_leaky_relu(shape, dtype, negative_slope):
999+
inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
1000+
ref_inp = to_reference(inp, True)
1001+
1002+
ref_out = torch.nn.functional.leaky_relu(ref_inp, negative_slope)
1003+
with flag_gems.use_gems():
1004+
res_out = torch.nn.functional.leaky_relu(inp, negative_slope)
1005+
1006+
gems_assert_close(res_out, ref_out, dtype)
1007+
1008+
1009+
@pytest.mark.leaky_relu_
1010+
@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
1011+
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
1012+
@pytest.mark.parametrize("negative_slope", [0.01, 0.2])
1013+
def test_leaky_relu_(shape, dtype, negative_slope):
1014+
inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
1015+
ref_inp = to_reference(inp.clone(), True)
1016+
1017+
ref_out = torch.nn.functional.leaky_relu(ref_inp, negative_slope, inplace=True)
1018+
with flag_gems.use_gems():
1019+
res_out = torch.nn.functional.leaky_relu(inp, negative_slope, inplace=True)
1020+
1021+
gems_assert_close(res_out, ref_out, dtype)
1022+
1023+
1024+
@pytest.mark.leaky_relu_out
1025+
@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
1026+
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
1027+
@pytest.mark.parametrize("negative_slope", [0.01, 0.2])
1028+
def test_leaky_relu_out(shape, dtype, negative_slope):
1029+
inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
1030+
ref_inp = to_reference(inp, True)
1031+
out = torch.empty_like(inp)
1032+
ref_out = torch.empty_like(ref_inp)
1033+
1034+
torch.ops.aten.leaky_relu.out(ref_inp, negative_slope, out=ref_out)
1035+
with flag_gems.use_gems():
1036+
torch.ops.aten.leaky_relu.out(inp, negative_slope, out=out)
1037+
1038+
gems_assert_close(out, ref_out, dtype)
1039+
1040+
1041+
@pytest.mark.leaky_relu_out
1042+
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
1043+
def test_leaky_relu_out_resizes_output(dtype):
1044+
inp = torch.randn((8, 16), dtype=dtype, device=flag_gems.device)
1045+
ref_inp = to_reference(inp, True)
1046+
out = torch.empty((32,), dtype=dtype, device=flag_gems.device)
1047+
ref_out = torch.empty((32,), dtype=dtype, device=ref_inp.device)
1048+
1049+
with warnings.catch_warnings():
1050+
warnings.simplefilter("ignore", UserWarning)
1051+
torch.ops.aten.leaky_relu.out(ref_inp, 0.01, out=ref_out)
1052+
with flag_gems.use_gems():
1053+
torch.ops.aten.leaky_relu.out(inp, 0.01, out=out)
1054+
1055+
assert out.shape == ref_out.shape
1056+
gems_assert_close(out, ref_out, dtype)
1057+
1058+
1059+
@pytest.mark.leaky_relu_out
1060+
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
1061+
def test_leaky_relu_out_dtype_mismatch(dtype):
1062+
inp = torch.randn((8, 16), dtype=dtype, device=flag_gems.device)
1063+
out_dtype = torch.float16 if dtype != torch.float16 else torch.float32
1064+
out = torch.empty_like(inp, dtype=out_dtype)
1065+
1066+
with flag_gems.use_gems(), pytest.raises(RuntimeError):
1067+
torch.ops.aten.leaky_relu.out(inp, 0.01, out=out)
1068+
1069+
9921070
@pytest.mark.relu6
9931071
@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
9941072
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)

0 commit comments

Comments
 (0)