Skip to content

Commit 7082124

Browse files
tengqmtengqm
committed
Split norm benchmarks
1 parent 754c602 commit 7082124

9 files changed

Lines changed: 376 additions & 258 deletions

‎benchmark/test_batch_norm.py‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
import pytest
2+
import torch
3+
4+
from . import attri_util as attr_utils
5+
from . import performance_utils as utils
6+
7+
8+
class NormBenchmark(utils.GenericBenchmark):
9+
# TODO: add new metric
10+
11+
def set_more_shapes(self):
12+
return [
13+
# 3D shapes represented as [batch_size, channels, hidden_size]
14+
(16, 16, 64),
15+
(16, 16, 1024),
16+
(16, 16, 4098),
17+
# 4D shapes represented as [batch_size, channels, H, W]
18+
(1, 8, 4, 4),
19+
(16, 8, 128, 128),
20+
]
21+
22+
23+
def input_fn(shape, dtype, device):
24+
C = shape[1]
25+
inp = torch.randn(shape, dtype=dtype, device=device)
26+
weight = torch.randn((C,), dtype=dtype, device=device)
27+
bias = torch.randn((C,), dtype=dtype, device=device)
28+
running_mean = None
29+
running_var = None
30+
training = True
31+
momentum = 0.1
32+
eps = 1e-5
33+
cudnn_enabled = True
34+
yield inp, weight, bias, running_mean, running_var, training, momentum, eps, cudnn_enabled
35+
36+
if utils.Config.bench_level == utils.BenchLevel.COMPREHENSIVE:
37+
running_mean = torch.randn((C,), dtype=dtype, device=device)
38+
running_var = torch.randn((C,), dtype=dtype, device=device)
39+
yield inp, weight, bias, running_mean, running_var, training, momentum, eps, cudnn_enabled
40+
41+
42+
@pytest.mark.batch_norm
43+
def test_batch_norm():
44+
bench = NormBenchmark(
45+
op_name="batch_norm",
46+
input_fn=input_fn,
47+
torch_op=torch.batch_norm,
48+
dtypes=attr_utils.FLOAT_DTYPES,
49+
)
50+
51+
bench.run()

‎benchmark/test_batch_norm_backward.py‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
import pytest
2+
import torch
3+
4+
import flag_gems
5+
6+
from . import attr_util as attr_utils
7+
from . import performance_utils as utils
8+
9+
10+
# TODO(Qiming): Consolidate this to a base package
11+
class NormBenchmark(utils.GenericBenchmark):
12+
# TODO: add new metric
13+
14+
def set_more_shapes(self):
15+
return [
16+
# 3D shapes represented as [batch_size, channels, hidden_size]
17+
(16, 16, 64),
18+
(16, 16, 1024),
19+
(16, 16, 4098),
20+
# 4D shapes represented as [batch_size, channels, H, W]
21+
(1, 8, 4, 4),
22+
(16, 8, 128, 128),
23+
]
24+
25+
26+
def batchnorm_input_fn(shape, dtype, device):
27+
C = shape[1]
28+
inp = torch.randn(shape, dtype=dtype, device=device)
29+
weight = torch.randn((C,), dtype=dtype, device=device)
30+
bias = torch.randn((C,), dtype=dtype, device=device)
31+
running_mean = None
32+
running_var = None
33+
training = True
34+
momentum = 0.1
35+
eps = 1e-5
36+
cudnn_enabled = True
37+
yield inp, weight, bias, running_mean, running_var, training, momentum, eps, cudnn_enabled
38+
39+
if utils.Config.bench_level == utils.BenchLevel.COMPREHENSIVE:
40+
running_mean = torch.randn((C,), dtype=dtype, device=device)
41+
running_var = torch.randn((C,), dtype=dtype, device=device)
42+
yield inp, weight, bias, running_mean, running_var, training, momentum, eps, cudnn_enabled
43+
44+
45+
@pytest.mark.batch_norm_backward
46+
def test_batch_norm_backward():
47+
def batch_norm_backward_input_fn(shape, dtype, device):
48+
for forward_args in batchnorm_input_fn(shape, dtype, device):
49+
(
50+
inp,
51+
weight,
52+
bias,
53+
running_mean,
54+
running_var,
55+
training,
56+
_,
57+
eps,
58+
_,
59+
) = forward_args
60+
61+
grad_output = torch.randn_like(inp)
62+
channels = weight.shape[0] if weight is not None else inp.shape[1]
63+
64+
if running_mean is None:
65+
running_mean = torch.zeros(channels, dtype=dtype, device=device)
66+
if running_var is None:
67+
running_var = torch.ones(channels, dtype=dtype, device=device)
68+
69+
save_mean = torch.randn(channels, dtype=torch.float32, device=device)
70+
save_invstd = torch.randn(channels, dtype=torch.float32, device=device)
71+
output_mask = [True, weight is not None, bias is not None]
72+
73+
yield (
74+
grad_output,
75+
inp,
76+
weight,
77+
running_mean,
78+
running_var,
79+
save_mean,
80+
save_invstd,
81+
training,
82+
eps,
83+
output_mask,
84+
)
85+
86+
bench = NormBenchmark(
87+
input_fn=batch_norm_backward_input_fn,
88+
op_name="native_batch_norm_backward",
89+
torch_op=torch.ops.aten.native_batch_norm_backward,
90+
dtypes=[torch.float32]
91+
if flag_gems.vendor_name == "mthreads"
92+
else attr_utils.FLOAT_DTYPES,
93+
)
94+
bench.set_gems(flag_gems.batch_norm_backward)
95+
96+
bench.run()

‎benchmark/test_group_norm.py‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
import pytest
2+
import torch
3+
4+
from . import attri_util as attr_utils
5+
from . import performance_utils as utils
6+
7+
8+
# TODO(Qiming): Extract this to a base class
9+
class NormBenchmark(utils.GenericBenchmark):
10+
# TODO: add new metric
11+
12+
def set_more_shapes(self):
13+
return [
14+
# 3D shapes represented as [batch_size, channels, hidden_size]
15+
(16, 16, 64),
16+
(16, 16, 1024),
17+
(16, 16, 4098),
18+
# 4D shapes represented as [batch_size, channels, H, W]
19+
(1, 8, 4, 4),
20+
(16, 8, 128, 128),
21+
]
22+
23+
24+
def group_norm_input_fn(shape, dtype, device):
25+
inp = torch.randn(shape, dtype=dtype, device=device)
26+
channel = shape[1]
27+
weight = torch.randn(
28+
[
29+
channel,
30+
],
31+
dtype=dtype,
32+
device=device,
33+
)
34+
bias = torch.randn(
35+
[
36+
channel,
37+
],
38+
dtype=dtype,
39+
device=device,
40+
)
41+
yield inp, channel // 2, weight, bias
42+
43+
if utils.Config.bench_level == utils.BenchLevel.COMPREHENSIVE:
44+
yield inp, channel, weight, bias
45+
46+
47+
@pytest.mark.group_norm
48+
def test_group_norm():
49+
bench = NormBenchmark(
50+
input_fn=group_norm_input_fn,
51+
op_name="group_norm",
52+
torch_op=torch.nn.functional.group_norm,
53+
dtypes=attr_utils.FLOAT_DTYPES,
54+
)
55+
bench.run()

‎benchmark/test_instance_norm.py‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
import pytest
2+
import torch
3+
4+
import flag_gems
5+
6+
from . import attri_util as attr_utils
7+
from . import performance_utils as utils
8+
9+
10+
class NormBenchmark(utils.GenericBenchmark):
11+
# TODO: add new metric
12+
13+
def set_more_shapes(self):
14+
return [
15+
# 3D shapes represented as [batch_size, channels, hidden_size]
16+
(16, 16, 64),
17+
(16, 16, 1024),
18+
(16, 16, 4098),
19+
# 4D shapes represented as [batch_size, channels, H, W]
20+
(1, 8, 4, 4),
21+
(16, 8, 128, 128),
22+
]
23+
24+
25+
def input_fn(shape, dtype, device):
26+
C = shape[1]
27+
inp = torch.randn(shape, dtype=dtype, device=device)
28+
weight = torch.randn((C,), dtype=dtype, device=device)
29+
bias = torch.randn((C,), dtype=dtype, device=device)
30+
running_mean = None
31+
running_var = None
32+
use_input_stats = True
33+
momentum = 0.1
34+
eps = 1e-5
35+
cudnn_enabled = True
36+
yield inp, weight, bias, running_mean, running_var, use_input_stats, momentum, eps, cudnn_enabled
37+
if utils.Config.bench_level == utils.BenchLevel.COMPREHENSIVE:
38+
running_mean = torch.randn((C,), dtype=dtype, device=device)
39+
running_var = torch.randn((C,), dtype=dtype, device=device)
40+
yield inp, weight, bias, running_mean, running_var, use_input_stats, momentum, eps, cudnn_enabled
41+
42+
43+
@pytest.mark.instance_norm
44+
def test_instance_norm(monkeypatch):
45+
if flag_gems.vendor_name == "kunlunxin" and utils.SkipVersion("torch", "<2.5"):
46+
pytest.skip(
47+
"BF16 is not supported in XPytorch 2.0. Please upgrade your PyTorch version >= 2.5"
48+
)
49+
50+
if flag_gems.vendor_name == "mthreads":
51+
# Compatible with older versions of LLVM
52+
monkeypatch.env("DISABLE_LLVM_OPT", "1")
53+
54+
bench = NormBenchmark(
55+
op_name="instance_norm",
56+
input_fn=input_fn,
57+
torch_op=torch.instance_norm,
58+
dtypes=attr_utils.FLOAT_DTYPES,
59+
)
60+
bench.set_gems(flag_gems.instance_norm)
61+
bench.run()

‎benchmark/test_layer_norm.py‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
import pytest
2+
import torch
3+
4+
from . import attri_util as attr_utils
5+
from . import performance_utils as utils
6+
7+
8+
# TODO(Qiming): Extract this to a base class
9+
class NormBenchmark(utils.GenericBenchmark):
10+
# TODO: add new metric
11+
12+
def set_more_shapes(self):
13+
return [
14+
# 3D shapes represented as [batch_size, channels, hidden_size]
15+
(16, 16, 64),
16+
(16, 16, 1024),
17+
(16, 16, 4098),
18+
# 4D shapes represented as [batch_size, channels, H, W]
19+
(1, 8, 4, 4),
20+
(16, 8, 128, 128),
21+
]
22+
23+
24+
def input_fn(shape, dtype, device):
25+
inp = torch.randn(shape, dtype=dtype, device=device)
26+
layer_shape = shape[1:]
27+
weight = torch.randn(layer_shape, dtype=dtype, device=device)
28+
bias = torch.randn(layer_shape, dtype=dtype, device=device)
29+
yield inp, layer_shape, weight, bias
30+
31+
32+
@pytest.mark.layer_norm
33+
def test_layer_norm():
34+
bench = NormBenchmark(
35+
op_name="layer_norm",
36+
input_fn=input_fn,
37+
torch_op=torch.layer_norm,
38+
dtypes=attr_utils.FLOAT_DTYPES,
39+
)
40+
bench.run()

0 commit comments

Comments
 (0)