cambricon: cambricon merge to master

Author: chenmiao1919

benchmark/test_special_perf.py

@@ -92,7 +92,7 @@ def isin_input_fn(shape, dtype, device):
         input_fn=isin_input_fn,
         op_name="isin",
         torch_op=torch.isin,
-        dtypes=[torch.int32] if vendor_name == "cambricon" else INT_DTYPES,
+        dtypes=INT_DTYPES,
     )
     bench.run()
 
@@ -108,7 +108,7 @@ def unique_input_fn(shape, dtype, device):
         input_fn=unique_input_fn,
         op_name="unique",
         torch_op=torch.unique,
-        dtypes=[torch.int32] if vendor_name == "cambricon" else INT_DTYPES,
+        dtypes=INT_DTYPES,
     )
     bench.run()
 
@@ -400,6 +400,7 @@ def diagonal_backward_input_fn(shape, dtype, device):
 
 @pytest.mark.skipif(flag_gems.device == "musa", reason="ZeroDivisionError")
 @pytest.mark.skipif(vendor_name == "kunlunxin", reason="RESULT TODOFIX")
+@pytest.mark.skipif(vendor_name == "cambricon", reason="TODOFIX")
 @pytest.mark.kron
 def test_perf_kron():
     class KronBenchmark(GenericBenchmark2DOnly):
@@ -428,8 +429,8 @@ def contiguous_input_fn(shape, dtype, device):
             inp = torch.randn(shape, dtype=dtype, device=device)
         else:
             inp = torch.randint(
-                low=-10000, high=10000, size=shape, dtype=dtype, device=device
-            )
+                low=-10000, high=10000, size=shape, dtype=dtype, device="cpu"
+            ).to(device)
         inp = inp[::2]
         yield inp,
 

benchmark/test_unary_pointwise_perf.py

@@ -14,6 +14,8 @@
 )
 from .performance_utils import Benchmark, generate_tensor_input
 
+fp64_is_supported = flag_gems.runtime.device.support_fp64
+
 
 class UnaryPointwiseBenchmark(Benchmark):
     """
@@ -136,7 +138,8 @@ def test_to_dtype_perf():
     bench = ToDtypeBenchmark(
         op_name="to",
         torch_op=torch.Tensor.to,
-        dtypes=[torch.float16, torch.bfloat16, torch.float64],
+        dtypes=[torch.float16, torch.bfloat16]
+        + ([torch.float64] if fp64_is_supported else []),
     )
     bench.run()
 

src/flag_gems/runtime/backend/_cambricon/__init__.py

@@ -29,9 +29,6 @@
     "sort",  # skip now
     "multinomial",  # skip now
     "_upsample_bicubic2d_aa",  # skip now
-    "batch_norm",  #
-    "pad",
-    "constant_pad_nd",  #
 )
 
 __all__ = ["*"]

src/flag_gems/runtime/backend/_cambricon/fused/__init__.py

@@ -1,11 +1,17 @@
+from .cross_entropy_loss import cross_entropy_loss
 from .fused_add_rms_norm import fused_add_rms_norm
 from .gelu_and_mul import gelu_and_mul
+from .outer import outer
 from .silu_and_mul import silu_and_mul
 from .skip_layernorm import skip_layer_norm
+from .weight_norm import weight_norm
 
 __all__ = [
     "skip_layer_norm",
     "fused_add_rms_norm",
     "silu_and_mul",
     "gelu_and_mul",
+    "cross_entropy_loss",
+    "outer",
+    "weight_norm",
 ]

…end/_cambricon/ops/cross_entropy_loss.py …d/_cambricon/fused/cross_entropy_loss.pysrc/flag_gems/runtime/backend/_cambricon/ops/cross_entropy_loss.py renamed to src/flag_gems/runtime/backend/_cambricon/fused/cross_entropy_loss.py src/flag_gems/runtime/backend/_cambricon/ops/cross_entropy_loss.py renamed to src/flag_gems/runtime/backend/_cambricon/fused/cross_entropy_loss.py

@@ -3,13 +3,14 @@
 import torch
 import triton
 import triton.language as tl
+from torch.nn import _reduction as _Reduction
 
 from flag_gems import runtime
 from flag_gems.runtime import torch_device_fn
 from flag_gems.utils import libentry
 
+from ..ops import sum
 from ..utils import TOTAL_CORE_NUM
-from .sum import sum
 
 logger = logging.getLogger(__name__)
 
@@ -999,8 +1000,13 @@ def backward(ctx, out_grad):
 
 
 def cross_entropy_loss(
-    inp, target, weight=None, reduction=1, ignore_index=-100, label_smoothing=0.0
+    inp, target, weight=None, reduction="mean", ignore_index=-100, label_smoothing=0.0
 ):
     return CrossEntropyLoss.apply(
-        inp, target, weight, reduction, ignore_index, label_smoothing
+        inp,
+        target,
+        weight,
+        _Reduction.get_enum(reduction),
+        ignore_index,
+        label_smoothing,
     )

src/flag_gems/runtime/backend/_cambricon/fused/fused_add_rms_norm.py

@@ -30,6 +30,7 @@ def get_configs():
 @triton.autotune(
     configs=get_configs(),
     key=["M", "N_COLS"],
+    restore_value=["x_ptr", "r_ptr"],
 )
 @libentry()
 @triton.jit(do_not_specialize=["eps"])
@@ -57,52 +58,37 @@ def fused_add_rms_norm_kernel(
         _mean = tl.zeros([M_BLOCK, BLOCK_SIZE], dtype=tl.float32)
         for offset in range(0, N_COLS, BLOCK_SIZE):
             cols = offset + tl.arange(0, BLOCK_SIZE)
-            mask = (m_offset < M)[:, None] & (cols < N_COLS)[None, :]
-            x = tl.load(
-                mx_ptr[:, None] + cols[None, :],
-                mask=mask,
-                other=0.0,
-                eviction_policy="evict_last",
-            ).to(tl.float32)
-            r = tl.load(
-                mr_ptr[:, None] + cols[None, :],
-                mask=mask,
-                other=0.0,
-                eviction_policy="evict_last",
-            ).to(tl.float32)
-            x += r
-            _mean += x * x
+            row_mask = m_offset < ub
+            col_mask = cols < N_COLS
+            mask = row_mask[:, None] & col_mask[None, :]
+            x = tl.load(mx_ptr[:, None] + cols[None, :], mask=mask, other=0.0).to(
+                tl.float32
+            )
+            r = tl.load(mr_ptr[:, None] + cols[None, :], mask=mask, other=0.0).to(
+                tl.float32
+            )
+            xpr = x + r
+            tl.store(mr_ptr[:, None] + cols[None, :], xpr, mask=mask)
+            _mean += xpr * xpr
 
         # Since `_mean * (1 / N_COLS)` performs better, make this change.
         # var = tl.sum(_mean / N_COLS, axis=1)
-        var = tl.sum(_mean * (1 / N_COLS), axis=1)
-        rrms = 1 / tl.sqrt(var + eps)
+        var = tl.sum(_mean * (1.0 / N_COLS), axis=1)
+        rrms = 1.0 / tl.sqrt(var + eps)
 
         for offset in range(0, N_COLS, BLOCK_SIZE):
             cols = offset + tl.arange(0, BLOCK_SIZE)
-            mask = (m_offset < M)[:, None] & (cols < N_COLS)[None, :]
-            x = tl.load(
-                mx_ptr[:, None] + cols[None, :],
-                mask=mask,
-                other=0.0,
-                eviction_policy="evict_first",
-            ).to(tl.float32)
-            r = tl.load(
-                mr_ptr[:, None] + cols[None, :],
-                mask=mask,
-                other=0.0,
-                eviction_policy="evict_last",
-            ).to(tl.float32)
-            x += r
-            w = tl.load(
-                w_ptr + cols,
-                mask=cols < N_COLS,
-                other=0.0,
-                eviction_policy="evict_first",
+            row_mask = m_offset < ub
+            col_mask = cols < N_COLS
+            mask = row_mask[:, None] & col_mask[None, :]
+
+            xpr = tl.load(mr_ptr[:, None] + cols[None, :], mask=mask, other=0.0).to(
+                tl.float32
             )
-            y = (x * rrms[:, None]).to(x_ptr.dtype.element_ty) * w
-            # write back to residual and input
-            tl.store(mr_ptr[:, None] + cols[None, :], x, mask=mask)
+            w = tl.load(w_ptr + cols, mask=col_mask, other=0.0).to(tl.float32)
+            y = xpr * rrms[:, None]
+            y = y * w
+            y = y.to(x_ptr.dtype.element_ty)
             tl.store(mx_ptr[:, None] + cols[None, :], y, mask=mask)
 
 

src/flag_gems/runtime/backend/_cambricon/fused/gelu_and_mul.py

@@ -21,6 +21,33 @@ def gelu_none_and_mul_kernel(x, y):
     return x_gelu * y
 
 
+@pointwise_dynamic(
+    promotion_methods=[(0, 1, 2, "DEFAULT"), (0, 1, 2, "DEFAULT")], num_outputs=2
+)
+@triton.jit
+def gelu_none_and_mul_grad_kernel(x, y, dgrad):
+    RCP_SQRT_2: tl.constexpr = 0.7071067811
+    COEFF: tl.constexpr = 0.7978845608028654
+
+    x_fp32 = x.to(tl.float32)
+    x_gelu = 0.5 * x_fp32 * (1 + fast_erf(x_fp32 * RCP_SQRT_2))
+
+    d_gelu = dgrad * y
+    dx = (
+        d_gelu
+        * 0.5
+        * (
+            1.0
+            + fast_erf(x_fp32 * RCP_SQRT_2)
+            + x_fp32 * COEFF * tl.exp(-0.5 * x_fp32 * x_fp32)
+        )
+    )
+
+    dy = dgrad * x_gelu
+
+    return dx, dy
+
+
 @pointwise_dynamic(promotion_methods=[(0, 1, "DEFAULT")])
 @triton.jit
 def gelu_tanh_and_mul_kernel(x, y):
@@ -33,17 +60,57 @@ def gelu_tanh_and_mul_kernel(x, y):
     return x_gelu * y
 
 
+@pointwise_dynamic(
+    promotion_methods=[(0, 1, 2, "DEFAULT"), (0, 1, 2, "DEFAULT")], num_outputs=2
+)
+@triton.jit
+def gelu_tanh_and_mul_grad_kernel(x, y, dgrad):
+    x_fp32 = x.to(tl.float32)
+    y_fp32 = y.to(tl.float32)
+
+    sqrt_2_over_pi = 0.7978845608028654  # sqrt(2 / pi)
+    a_cubed = x_fp32 * x_fp32 * x_fp32
+    tanh_arg = sqrt_2_over_pi * (x_fp32 + 0.044715 * a_cubed)
+    tanh_result = fast_tanh(tanh_arg)
+    geglu_a = 0.5 * x_fp32 * (1 + tanh_result)
+    dy = geglu_a * dgrad
+
+    term1 = 0.5 * (1 + tanh_result)
+    tanh_sq = tanh_result * tanh_result
+    term2 = (
+        0.5
+        * x_fp32
+        * (1 - tanh_sq)
+        * (sqrt_2_over_pi * (1 + 3 * 0.044715 * x_fp32 * x_fp32))
+    )
+    dx = dgrad * y_fp32 * (term1 + term2)
+
+    return dx, dy
+
+
 class GeluAndMul(torch.autograd.Function):
     @staticmethod
-    def forward(ctx, A, B, approximate="none"):
+    def forward(ctx, x, y, approximate="none"):
         logger.debug("GEMS_CAMBRICON GELU AND MUL FORWARD")
+        ctx.save_for_backward(x, y)
+        ctx.approximate = approximate
         if approximate == "none":
-            return gelu_none_and_mul_kernel(A, B)
+            return gelu_none_and_mul_kernel(x, y)
         elif approximate == "tanh":
-            return gelu_tanh_and_mul_kernel(A, B)
+            return gelu_tanh_and_mul_kernel(x, y)
         else:
             raise ValueError(f"Invalid approximate value: {approximate}")
 
+    @staticmethod
+    def backward(ctx, dgrad):
+        logging.debug("GEMS GELU AND MUL BACKWARD")
+        x, y = ctx.saved_tensors
+        if ctx.approximate == "none":
+            dx, dy = gelu_none_and_mul_grad_kernel(x, y, dgrad)
+        else:
+            dx, dy = gelu_tanh_and_mul_grad_kernel(x, y, dgrad)
+        return dx, dy, None
+
 
 def gelu_and_mul(A, B, approximate="none"):
     return GeluAndMul.apply(A, B, approximate)

…/runtime/backend/_cambricon/ops/outer.py …untime/backend/_cambricon/fused/outer.pysrc/flag_gems/runtime/backend/_cambricon/ops/outer.py renamed to src/flag_gems/runtime/backend/_cambricon/fused/outer.py src/flag_gems/runtime/backend/_cambricon/ops/outer.py renamed to src/flag_gems/runtime/backend/_cambricon/fused/outer.py

@@ -6,8 +6,8 @@
 
 from flag_gems.utils import libentry
 
+from ..ops import mv
 from ..utils import TOTAL_CORE_NUM
-from .mv import mv
 
 logger = logging.getLogger(__name__)
 

src/flag_gems/runtime/backend/_cambricon/fused/silu_and_mul.py

@@ -17,12 +17,32 @@ def silu_and_mul_kernel(x, y):
     return x_silu * y
 
 
+@pointwise_dynamic(
+    promotion_methods=[(0, 1, 2, "DEFAULT"), (0, 1, 2, "DEFAULT")], num_outputs=2
+)
+@triton.jit
+def silu_and_mul_grad_kernel(x, y, dgrad):
+    x_fp32 = x.to(tl.float32)
+    sig = tl.extra.mlu.libdevice.fast_sigmoid(x_fp32)
+    x_silu = x_fp32 * sig
+    d_x_silu = sig * (1 + x_fp32 * (1 - sig))
+    dx = d_x_silu * dgrad * y
+    dy = dgrad * x_silu
+    return dx, dy
+
+
 class SiluAndMul(torch.autograd.Function):
     @staticmethod
     def forward(ctx, A, B):
+        ctx.save_for_backward(A, B)
         logger.debug("GEMS_CAMBRICON SILU AND MUL FORWARD")
         return silu_and_mul_kernel(A, B)
 
+    def backward(ctx, grad_output):
+        A, B = ctx.saved_tensors
+        grad_A, grad_B = silu_and_mul_grad_kernel(A, B, grad_output)
+        return grad_A, grad_B
+
 
 def silu_and_mul(A, B):
     return SiluAndMul.apply(A, B)