Merge branch 'master' into pr/_fake_quantize_learnable_per_channel_affine_backward

Author: bin913

benchmark/test_conj_physical.py

@@ -34,7 +34,6 @@ def _input_fn(shape, dtype, device):
 
 
 class Conj_physicalBenchmark(base.GenericBenchmarkExcluse3D):
-    # TODO(Qiming): Check if this is necessary
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
@@ -55,7 +54,14 @@ def set_more_shapes(self):
 
 @pytest.mark.conj_physical
 def test_conj_physical():
-    dtypes = consts.FLOAT_DTYPES + consts.INT_DTYPES + consts.COMPLEX_DTYPES
+    if "npu" in flag_gems.device or "ascend" in flag_gems.device.lower():
+        # Ascend NPU: kernel mode event timing is unstable, use operator mode
+        from .conftest import Config
+
+        Config.mode = consts.BenchMode.OPERATOR
+        dtypes = consts.FLOAT_DTYPES + consts.INT_DTYPES
+    else:
+        dtypes = consts.FLOAT_DTYPES + consts.INT_DTYPES + consts.COMPLEX_DTYPES
 
     bench = Conj_physicalBenchmark(
         input_fn=_input_fn,

benchmark/test_layer_norm.py

@@ -31,6 +31,12 @@ def set_more_shapes(self):
             # 4D shapes represented as [batch_size, channels, H, W]
             (1, 8, 4, 4),
             (16, 8, 128, 128),
+            # Medium normalized dimensions across representative row counts
+            (256, 512),
+            (4096, 256),
+            (4096, 512),
+            (2048, 1024),
+            (1024, 2048),
         ]
 
 

benchmark/test_linalg_det.py

@@ -0,0 +1,123 @@
+import math
+
+import pytest
+import torch
+
+import flag_gems
+from flag_gems import linalg_det, linalg_det_out
+
+from . import base, consts
+from .conftest import Config
+
+VENDOR = flag_gems.vendor_name
+
+
+if VENDOR == "ascend":
+    Config.mode = consts.BenchMode.OPERATOR
+
+
+def _small_ops_det(A):
+    n = A.shape[-1]
+    batch_shape = A.shape[:-2]
+    B = math.prod(batch_shape) if batch_shape else 1
+    LU = A.clone().reshape(B, n, n)
+    sign = torch.ones(B, dtype=A.dtype, device=A.device)
+    bidx = torch.arange(B, device=A.device)
+    for k in range(n):
+        p = LU[:, k:, k].abs().argmax(dim=-1) + k
+        swap = p != k
+        sign = torch.where(swap, -sign, sign)
+        row_k = LU[:, k, :].clone()
+        row_p = LU[bidx, p, :].clone()
+        LU[:, k, :] = row_p
+        LU[bidx[swap], p[swap], :] = row_k[swap]
+        pivot = LU[:, k, k]
+        safe_pivot = torch.where(pivot == 0, torch.ones_like(pivot), pivot)
+        col = LU[:, k + 1 :, k]
+        mult = torch.where(
+            (pivot == 0).unsqueeze(-1),
+            torch.zeros_like(col),
+            col / safe_pivot.unsqueeze(-1),
+        )
+        LU[:, k + 1 :, k] = mult
+        LU[:, k + 1 :, k + 1 :] -= mult.unsqueeze(-1) * LU[:, k : k + 1, k + 1 :]
+    det = LU.diagonal(dim1=-2, dim2=-1).prod(dim=-1) * sign
+    return det.reshape(batch_shape)
+
+
+def _torch_det(A):
+    if A.device.type == "npu":
+        return _small_ops_det(A)
+    return torch.linalg.det(A)
+
+
+def _torch_det_out(A, *, out):
+    if A.device.type == "npu":
+        out.copy_(_torch_det(A))
+        return out
+    return torch.linalg.det(A, out=out)
+
+
+DET_SHAPES = [
+    (16, 16),
+    (32, 32),
+    (64, 64),
+    (128, 128),
+    (256, 256),
+    (4096, 4, 4),
+    (1024, 8, 8),
+    (1024, 16, 16),
+    (128, 16, 16),
+    (4, 32, 32),
+    (512, 32, 32),
+    (256, 64, 64),
+    (32, 128, 128),
+    (8, 256, 256),
+]
+
+DET_DTYPES = [torch.float32] + (
+    [torch.float64] if flag_gems.runtime.device.support_fp64 else []
+)
+
+
+class DetBenchmark(base.Benchmark):
+    def set_shapes(self, shape_file_path=None):
+        self.shapes = DET_SHAPES
+
+    def get_input_iter(self, cur_dtype):
+        for shape in self.shapes:
+            A = torch.randn(shape, dtype=cur_dtype, device=self.device)
+            yield (A,)
+
+
+@pytest.mark.linalg_det
+def test_linalg_det():
+    bench = DetBenchmark(
+        op_name="linalg_det",
+        torch_op=_torch_det,
+        dtypes=DET_DTYPES,
+    )
+    bench.set_gems(linalg_det)
+    bench.run()
+
+
+class DetOutBenchmark(base.Benchmark):
+    def set_shapes(self, shape_file_path=None):
+        self.shapes = DET_SHAPES
+
+    def get_input_iter(self, cur_dtype):
+        for shape in self.shapes:
+            A = torch.randn(shape, dtype=cur_dtype, device=self.device)
+            out = torch.empty(shape[:-2], dtype=cur_dtype, device=self.device)
+            yield (A, {"out": out})
+
+
+@pytest.mark.linalg_det_out
+def test_linalg_det_out():
+    bench = DetOutBenchmark(
+        op_name="linalg_det_out",
+        torch_op=_torch_det_out,
+        dtypes=DET_DTYPES,
+    )
+    bench.set_gems(linalg_det_out)
+    bench.run()

benchmark/test_replication_pad2d_backward.py

@@ -1,7 +1,15 @@
 import pytest
 import torch
 
+import flag_gems
+
 from . import base, consts
+from .conftest import Config
+
+VENDOR = flag_gems.vendor_name
+
+if VENDOR == "ascend":
+    Config.mode = consts.BenchMode.OPERATOR
 
 REPLICATION_PAD2D_BACKWARD_SHAPES = [
     (1, 3, 256, 256),

conf/operators.yaml

@@ -5770,6 +5770,26 @@ ops:
       - aten
     kind:
       - Math
+  - id: linalg_det
+    description: |
+      Computes the determinant of a square matrix via LU decomposition with partial pivoting.
+    for:
+      - linalg.det
+    labels:
+      - aten
+    kind:
+      - LinearAlg
+    stages:
+      - alpha: '5.4'
+  - id: linalg_det_out
+    description: |
+      A variant of linalg_det that writes the determinant to the provided out tensor.
+    for:
+      - linalg.det.out
+    labels:
+      - aten
+    kind:
+      - LinearAlg
     stages:
       - alpha: '5.4'
   - id: linalg_lu_factor_ex

src/flag_gems/__init__.py

@@ -629,6 +629,8 @@ def torch_ge(v):
     ("linalg_cholesky", linalg_cholesky),
     ("linalg_cross", linalg_cross),
     ("linalg_cross.out", linalg_cross_out),
+    ("linalg_det", linalg_det),
+    ("linalg_det.out", linalg_det_out),
     ("linalg_ldl_factor", ldl_factor),
     ("linalg_ldl_factor_ex", ldl_factor_ex),
     ("linalg_lu_factor", linalg_lu_factor),

src/flag_gems/ops/__init__.py

@@ -436,6 +436,7 @@
 from flag_gems.ops.lift_fresh_copy import lift_fresh_copy, lift_fresh_copy_out
 from flag_gems.ops.linalg_cholesky import linalg_cholesky
 from flag_gems.ops.linalg_cross import linalg_cross, linalg_cross_out
+from flag_gems.ops.linalg_det import linalg_det, linalg_det_out
 from flag_gems.ops.linalg_ldl_factor import ldl_factor
 from flag_gems.ops.linalg_ldl_solve import linalg_ldl_solve
 from flag_gems.ops.linalg_lstsq import linalg_lstsq
@@ -1309,6 +1310,8 @@
     "linalg_cholesky",
     "linalg_cross",
     "linalg_cross_out",
+    "linalg_det",
+    "linalg_det_out",
     "linalg_ldl_solve",
     "linalg_lstsq",
     "linalg_lu_factor",