Merge remote-tracking branch 'upstream/master' into pr/_fake_quantize_learnable_per_channel_affine_backward

Author: chx7514

benchmark/test__batch_norm_impl_index.py

@@ -0,0 +1,40 @@
+import pytest
+import torch
+
+from . import base, consts
+
+
+class NormBenchmark(base.GenericBenchmark):
+    def set_more_shapes(self):
+        return [
+            # 3D shapes represented as [batch_size, channels, hidden_size]
+            (16, 16, 64),
+            (16, 16, 1024),
+            (16, 16, 4098),
+            # 4D shapes represented as [batch_size, channels, H, W]
+            (1, 8, 4, 4),
+            (16, 8, 128, 128),
+        ]
+
+
+@pytest.mark.batch_norm_impl_index
+def test__batch_norm_impl_index():
+    def batch_norm_impl_index_input_fn(shape, dtype, device):
+        C = shape[1]
+        inp = torch.randn(shape, dtype=dtype, device=device)
+        weight = torch.randn(C, dtype=dtype, device=device)
+        bias = torch.randn(C, dtype=dtype, device=device)
+        running_mean = torch.zeros(C, dtype=dtype, device=device)
+        running_var = torch.ones(C, dtype=dtype, device=device)
+        yield inp, weight, bias, running_mean, running_var, True, 0.1, 1e-5, True
+
+    bench = NormBenchmark(
+        input_fn=batch_norm_impl_index_input_fn,
+        op_name="_batch_norm_impl_index",
+        torch_op=torch._batch_norm_impl_index,
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    from flag_gems.ops._batch_norm_impl_index import batch_norm_impl_index as gems_bn
+
+    bench.set_gems(gems_bn)
+    bench.run()

benchmark/test_adaptive_avg_pool3d_backward.py

@@ -0,0 +1,36 @@
+import pytest
+import torch
+
+from . import base, consts
+
+# Shapes for adaptive_avg_pool3d backward benchmark
+ADAPTIVE_AVG_POOL3D_BACKWARD_SHAPES = [
+    (1, 3, 8, 8, 8),
+    (2, 3, 16, 16, 16),
+    (1, 1, 32, 32, 32),
+    (4, 8, 64, 64, 64),
+]
+
+
+class AdaptiveAvgPool3DBackwardBenchmark(base.Benchmark):
+    def set_shapes(self, shape_file_path=None):
+        self.shapes = ADAPTIVE_AVG_POOL3D_BACKWARD_SHAPES
+        self.output_sizes = [(4, 4, 4), (8, 8, 8), (16, 16, 16), (32, 32, 32)]
+
+    def get_input_iter(self, cur_dtype):
+        for shape, output_size in zip(self.shapes, self.output_sizes):
+            x = torch.randn(shape, dtype=cur_dtype, device=self.device)
+            # Compute forward to get output shape
+            out = torch.nn.functional.adaptive_avg_pool3d(x, output_size)
+            grad = torch.ones_like(out)
+            yield grad, x
+
+
+@pytest.mark.adaptive_avg_pool3d_backward
+def test_adaptive_avg_pool3d_backward():
+    bench = AdaptiveAvgPool3DBackwardBenchmark(
+        op_name="adaptive_avg_pool3d_backward",
+        torch_op=torch.ops.aten._adaptive_avg_pool3d_backward,
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.run()

benchmark/test_addbmm_.py

@@ -0,0 +1,30 @@
+import pytest
+import torch
+
+from . import base, consts
+
+
+def _input_fn(b, m, n, k, dtype, device, b_column_major):
+    inp1 = torch.randn([b, m, k], dtype=dtype, device=device, requires_grad=True)
+
+    if b_column_major:
+        inp2 = torch.randn([b, n, k], dtype=dtype, device=device, requires_grad=True)
+        inp2 = inp2.transpose(1, 2).contiguous()
+    else:
+        inp2 = torch.randn([b, k, n], dtype=dtype, device=device, requires_grad=True)
+
+    # addbmm_ is inplace, bias must not require grad (leaf tensor constraint)
+    bias = torch.randn([m, n], dtype=dtype, device=device, requires_grad=False)
+
+    yield bias, inp1, inp2
+
+
+@pytest.mark.addbmm_
+def test_addbmm_():
+    bench = base.BlasBenchmark(
+        op_name="addbmm_",
+        input_fn=_input_fn,
+        torch_op=lambda bias, inp1, inp2: bias.addbmm_(inp1, inp2),
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.run()

benchmark/test_addr_.py

@@ -0,0 +1,35 @@
+# Generated by KernelGen: https://github.com/flagos-ai/KernelGen
+from typing import Generator
+
+import pytest
+import torch
+
+from . import base, consts
+
+
+class AddrInplaceBenchmark(base.BlasBenchmark):
+    def set_more_shapes(self):
+        return []
+
+    def get_input_iter(self, dtype) -> Generator:
+        for shape in self.shapes:
+            m, n = shape[0], shape[1]
+            yield from self.input_fn(m, n, dtype, self.device)
+
+
+def _input_fn(m, n, cur_dtype, device):
+    inp1 = torch.randn([m, n], dtype=cur_dtype, device=device)
+    inp2 = torch.randn([m], dtype=cur_dtype, device=device)
+    inp3 = torch.randn([n], dtype=cur_dtype, device=device)
+    yield inp1, inp2, inp3, {"alpha": 0.5, "beta": 0.5}
+
+
+@pytest.mark.addr_
+def test_addr_():
+    bench = AddrInplaceBenchmark(
+        op_name="addr_",
+        input_fn=_input_fn,
+        torch_op=torch.Tensor.addr_,
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.run()

benchmark/test_avg_pool1d.py

@@ -0,0 +1,77 @@
+# Copyright 2026 FlagOS Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Generator
+
+import pytest
+import torch
+
+from . import base, consts, utils
+
+
+class AvgPool1dBenchmark(base.GenericBenchmark):
+    def get_input_iter(self, dtype) -> Generator:
+        shapes_3d = [
+            (4, 3, 224),
+            (16, 64, 128),
+            (32, 128, 64),
+            (64, 256, 32),
+            (128, 512, 16),
+        ]
+
+        for shape in shapes_3d:
+            yield from self.input_fn(shape, dtype, self.device)
+
+
+def avg_pool1d_input_fn(shape, dtype, device):
+    inp = utils.generate_tensor_input(shape, dtype, device)
+
+    # Common case
+    yield inp, {
+        "kernel_size": [3],
+        "stride": [2],
+        "padding": [1],
+        "ceil_mode": False,
+        "count_include_pad": True,
+    }
+
+    if base.Config.bench_level == consts.BenchLevel.COMPREHENSIVE:
+        # With count_include_pad=False
+        yield inp, {
+            "kernel_size": [3],
+            "stride": [2],
+            "padding": [1],
+            "ceil_mode": False,
+            "count_include_pad": False,
+        }
+
+        # With ceil_mode
+        yield inp, {
+            "kernel_size": [3],
+            "stride": [2],
+            "padding": [1],
+            "ceil_mode": True,
+            "count_include_pad": True,
+        }
+
+
+@pytest.mark.avg_pool1d
+def test_avg_pool1d():
+    bench = AvgPool1dBenchmark(
+        input_fn=avg_pool1d_input_fn,
+        op_name="avg_pool1d",
+        torch_op=torch.ops.aten.avg_pool1d,
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.run()

benchmark/test_cdist.py

@@ -0,0 +1,41 @@
+import pytest
+import torch
+
+from . import base
+
+# Shapes for cdist benchmark: (P, M), (R, M) -> (P, R)
+# torch.cdist doesn't support float16 on CUDA
+CDIST_SHAPES = [
+    ((4, 8), (6, 8)),
+    ((8, 16), (8, 16)),
+    ((16, 32), (16, 32)),
+    ((32, 64), (32, 64)),
+    ((64, 128), (64, 128)),
+]
+
+
+class CdistBenchmark(base.Benchmark):
+    def set_shapes(self, shape_file_path=None):
+        self.shapes = CDIST_SHAPES
+
+    def get_input_iter(self, cur_dtype):
+        for shape1, shape2 in self.shapes:
+            x1 = torch.randn(*shape1, dtype=cur_dtype, device=self.device)
+            x2 = torch.randn(*shape2, dtype=cur_dtype, device=self.device)
+            yield x1, x2, 2.0
+
+    def get_tflops(self, op, *args, **kwargs):
+        x1, x2, _ = args
+        # FLOPs = 2 * P * R * M (for L2 distance computation)
+        return 2 * x1.shape[-2] * x2.shape[-2] * x1.shape[-1]
+
+
+@pytest.mark.cdist
+def test_cdist():
+    bench = CdistBenchmark(
+        op_name="cdist",
+        torch_op=torch.cdist,
+        # torch.cdist doesn't support float16 on CUDA; only float32 is numerically stable
+        dtypes=[torch.float32],
+    )
+    bench.run()

benchmark/test_dropout.py

@@ -38,13 +38,13 @@ def test_dropout():
     bench.run()
 
 
-@pytest.mark.dropout_backward
+@pytest.mark.native_dropout_backward
 @pytest.mark.skipif(
     flag_gems.vendor_name == "tsingmicro", reason="Issue #4131: not working"
 )
-def test_dropout_backward():
+def test_native_dropout_backward():
     bench = base.GenericBenchmark(
-        op_name="dropout_backward",
+        op_name="native_dropout_backward",
         input_fn=_dropout_backward_input_fn,
         torch_op=torch.ops.aten.native_dropout_backward,
         dtypes=consts.FLOAT_DTYPES,

benchmark/test_empty.py

@@ -22,6 +22,23 @@ def empty_input_fn(shape, dtype, device):
     yield shape
 
 
+def empty_permuted_input_fn(shape, dtype, device):
+    # Reverse the physical layout so the allocation exercises a non-contiguous
+    # memory ordering rather than the plain contiguous one.
+    yield shape, list(reversed(range(len(shape))))
+
+
+@pytest.mark.empty_permuted
+def test_empty_permuted():
+    bench = base.GenericBenchmark(
+        op_name="empty_permuted",
+        torch_op=torch.empty_permuted,
+        dtypes=consts.FLOAT_DTYPES,
+        input_fn=empty_permuted_input_fn,
+    )
+    bench.run()
+
+
 @pytest.mark.empty
 def test_empty():
     bench = base.GenericBenchmark(

benchmark/test_fake_quantize_per_channel_affine.py

@@ -0,0 +1,70 @@
+# Copyright 2026 FlagOS Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+from . import base, consts
+
+
+@pytest.mark.fake_quantize_per_channel_affine
+def test_fake_quantize_per_channel_affine():
+    class BenchmarkFakeQuantizePerChannelAffine(base.Benchmark):
+        """
+        Benchmark fake_quantize_per_channel_affine operator
+        """
+
+        axis_configs = (0, 1)
+        DEFAULT_SHAPES = [
+            (4, 4),
+            (64, 64),
+            (128, 256),
+            (512, 512),
+            (1024, 1024),
+            (2, 3, 128, 128),
+            (8, 16, 64, 64),
+        ]
+
+        def set_shapes(self, shape_file_path=None):
+            self.shapes = self.DEFAULT_SHAPES
+
+        def get_input_iter(self, dtype):
+            for shape in self.shapes:
+                for axis in self.axis_configs:
+                    if axis >= len(shape):
+                        continue
+                    inp = torch.randn(shape, dtype=dtype, device="cuda")
+                    n_channels = shape[axis]
+                    scale = (
+                        torch.rand(n_channels, dtype=torch.float32, device="cuda") * 0.1
+                        + 0.01
+                    )
+                    zero_point = torch.zeros(
+                        n_channels, dtype=torch.int32, device="cuda"
+                    )
+                    quant_min = 0
+                    quant_max = 255
+                    yield inp, scale, zero_point, axis, quant_min, quant_max
+
+        def forward(self, inp, scale, zero_point, axis, quant_min, quant_max):
+            return torch.fake_quantize_per_channel_affine(
+                inp, scale, zero_point, axis, quant_min, quant_max
+            )
+
+    bench = BenchmarkFakeQuantizePerChannelAffine(
+        op_name="fake_quantize_per_channel_affine",
+        torch_op=torch.fake_quantize_per_channel_affine,
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.run()