[KernelGen] Support bfloat16 for bitwise_not bitwise_or div embedding logical_and polar on Ascend and update tests & benchmark

Author: ftgreat

benchmark/test_binary_pointwise_perf.py

@@ -7,6 +7,27 @@
 from benchmark.performance_utils import Benchmark, generate_tensor_input
 
 
+def _polar_baseline(abs, angle):
+    if abs.dtype == torch.bfloat16:
+        return torch.polar(abs.to(torch.float32), angle.to(torch.float32))
+    return torch.polar(abs, angle)
+
+
+def _bitwise_or_baseline(a, b):
+    if a.dtype == torch.bfloat16:
+        return torch.bitwise_or(a.view(torch.int16), b.view(torch.int16)).view(
+            torch.bfloat16
+        )
+    return torch.bitwise_or(a, b)
+
+
+def _bitwise_or_inplace_baseline(a, b):
+    if a.dtype == torch.bfloat16:
+        a.view(torch.int16).bitwise_or_(b.view(torch.int16))
+        return a
+    return a.bitwise_or_(b)
+
+
 class BinaryPointwiseBenchmark(Benchmark):
     """
     Base class for benchmarking binary pointwise operations.
@@ -47,11 +68,11 @@ def get_tflops(self, op, *args, **kwargs):
             ("mul", torch.mul, FLOAT_DTYPES),
             ("sub", torch.sub, FLOAT_DTYPES),
             ("pow", torch.pow, FLOAT_DTYPES),
-            ("polar", torch.polar, [torch.float32]),
+            ("polar", _polar_baseline, [torch.float32, torch.bfloat16]),
             ("floor_divide", torch.floor_divide, INT_DTYPES),
             ("remainder", torch.remainder, INT_DTYPES),
             ("logical_or", torch.logical_or, INT_DTYPES + BOOL_DTYPES),
-            ("logical_and", torch.logical_and, INT_DTYPES + BOOL_DTYPES),
+            ("logical_and", torch.logical_and, FLOAT_DTYPES + INT_DTYPES + BOOL_DTYPES),
             ("logical_xor", torch.logical_xor, INT_DTYPES + BOOL_DTYPES),
             # Comparison operations
             ("eq", torch.eq, FLOAT_DTYPES),
@@ -66,7 +87,7 @@ def get_tflops(self, op, *args, **kwargs):
             ("minimum", torch.minimum, FLOAT_DTYPES),
             # Bitwise operations
             ("bitwise_and", torch.bitwise_and, INT_DTYPES + BOOL_DTYPES),
-            ("bitwise_or", torch.bitwise_or, INT_DTYPES + BOOL_DTYPES),
+            ("bitwise_or", _bitwise_or_baseline, INT_DTYPES + BOOL_DTYPES + [torch.bfloat16]),
             # Numerical Checks
             ("isclose", torch.isclose, FLOAT_DTYPES + INT_DTYPES),
             ("allclose", torch.allclose, FLOAT_DTYPES + INT_DTYPES),
@@ -97,10 +118,10 @@ def test_general_binary_pointwise_perf(op_name, torch_op, dtypes):
             ("floor_divide_", lambda a, b: a.floor_divide_(b), INT_DTYPES),
             ("remainder_", lambda a, b: a.remainder_(b), INT_DTYPES),
             ("logical_or_", lambda a, b: a.logical_or_(b), INT_DTYPES + BOOL_DTYPES),
-            ("logical_and_", lambda a, b: a.logical_and_(b), INT_DTYPES + BOOL_DTYPES),
+            ("logical_and_", lambda a, b: a.logical_and_(b), FLOAT_DTYPES + INT_DTYPES + BOOL_DTYPES),
             # Bitwise operations
             ("bitwise_and_", lambda a, b: a.bitwise_and_(b), INT_DTYPES + BOOL_DTYPES),
-            ("bitwise_or_", lambda a, b: a.bitwise_or_(b), INT_DTYPES + BOOL_DTYPES),
+            ("bitwise_or_", _bitwise_or_inplace_baseline, INT_DTYPES + BOOL_DTYPES + [torch.bfloat16]),
         ]
     ],
 )

benchmark/test_special_perf.py

@@ -365,10 +365,7 @@ def test_perf_embedding():
         input_fn=embedding_input_fn,
         op_name="embedding",
         torch_op=torch.nn.functional.embedding,
-        dtypes=[
-            torch.float32,
-            torch.float16,
-        ],  # Note(Zhengzekang): triton do not support bfloat16 atomic add which is used in embedding grad.
+        dtypes=FLOAT_DTYPES,
     )
     bench.run()
 
@@ -379,10 +376,7 @@ def test_perf_embedding_backward():
         input_fn=embedding_backward_input_fn,
         op_name="embedding",
         torch_op=torch.nn.functional.embedding,
-        dtypes=[
-            torch.float32,
-            torch.float16,
-        ],  # Note(Zhengzekang): triton do not support bfloat16 atomic add which is used in embedding grad.
+        dtypes=FLOAT_DTYPES,
         is_backward=True,
     )
     bench.run()

benchmark/test_unary_pointwise_perf.py

@@ -17,6 +17,19 @@
 fp64_is_supported = flag_gems.runtime.device.support_fp64
 
 
+def _bitwise_not_baseline(a):
+    if a.dtype == torch.bfloat16:
+        return torch.bitwise_not(a.view(torch.int16)).view(torch.bfloat16)
+    return torch.bitwise_not(a)
+
+
+def _bitwise_not_inplace_baseline(a):
+    if a.dtype == torch.bfloat16:
+        a.view(torch.int16).bitwise_not_()
+        return a
+    return a.bitwise_not_()
+
+
 class UnaryPointwiseBenchmark(Benchmark):
     """
     Base class for benchmarking unary pointwise operations.
@@ -72,7 +85,7 @@ def get_tflops(self, op, *args, **kwargs):
     ("atan", torch.atan, FLOAT_DTYPES),
     ("acos", torch.acos, FLOAT_DTYPES),
     # Bitwise operations
-    ("bitwise_not", torch.bitwise_not, INT_DTYPES),
+    ("bitwise_not", _bitwise_not_baseline, INT_DTYPES + [torch.bfloat16]),
     # Numerical Checks
     ("isinf", torch.isinf, FLOAT_DTYPES),
     ("isnan", torch.isnan, FLOAT_DTYPES),
@@ -127,7 +140,7 @@ def test_general_unary_pointwise_perf(op_name, torch_op, dtypes):
     ("tanh_", torch.tanh_, FLOAT_DTYPES),
     ("atan_", torch.atan_, FLOAT_DTYPES),
     # Bitwise operations
-    ("bitwise_not_", lambda a: a.bitwise_not_(), INT_DTYPES),
+    ("bitwise_not_", _bitwise_not_inplace_baseline, INT_DTYPES + [torch.bfloat16]),
 ]
 
 

src/flag_gems/runtime/backend/_ascend/ops/__init__.py

@@ -2,6 +2,14 @@
 from .all import all, all_dim, all_dims
 from .amax import amax
 from .angle import angle
+from .bitwise_not import bitwise_not, bitwise_not_
+from .bitwise_or import (
+    bitwise_or_scalar,
+    bitwise_or_scalar_,
+    bitwise_or_scalar_tensor,
+    bitwise_or_tensor,
+    bitwise_or_tensor_,
+)
 from .any import any, any_dim, any_dims
 from .arange import arange
 from .argmax import argmax
@@ -10,6 +18,17 @@
 from .cat import cat
 from .count_nonzero import count_nonzero
 from .cumsum import cumsum, normed_cumsum
+from .div import (
+    div_mode,
+    div_mode_,
+    floor_divide,
+    floor_divide_,
+    true_divide,
+    true_divide_,
+    true_divide_out,
+    trunc_divide,
+    trunc_divide_,
+)
 from .diag import diag
 from .diag_embed import diag_embed
 from .diagonal import diagonal_backward
@@ -28,6 +47,7 @@
 from .index_select import index_select
 from .isin import isin
 from .linspace import linspace
+from .logical_and import logical_and, logical_and_
 from .log_softmax import log_softmax, log_softmax_backward
 from .masked_fill import masked_fill, masked_fill_
 from .masked_select import masked_select
@@ -148,9 +168,27 @@
     "any_dims",
     "any_dim",
     "angle",
+    "bitwise_not",
+    "bitwise_not_",
+    "bitwise_or_scalar",
+    "bitwise_or_scalar_",
+    "bitwise_or_scalar_tensor",
+    "bitwise_or_tensor",
+    "bitwise_or_tensor_",
     "multinomial",
     "index_add",
     "_unique2",
     "upsample_nearest2d",
     "randperm",
+    "true_divide",
+    "true_divide_",
+    "true_divide_out",
+    "trunc_divide",
+    "trunc_divide_",
+    "floor_divide",
+    "floor_divide_",
+    "div_mode",
+    "div_mode_",
+    "logical_and",
+    "logical_and_",
 ]

src/flag_gems/runtime/backend/_ascend/ops/bitwise_not.py

@@ -0,0 +1,30 @@
+import logging
+
+import torch
+import triton
+
+from flag_gems.utils import pointwise_dynamic
+
+logger = logging.getLogger(f'flag_gems.runtime._ascend.ops.{__name__.split(".")[-1]}')
+
+
+@pointwise_dynamic(promotion_methods=[(0, "DEFAULT")])
+@triton.jit
+def bitwise_not_func(x):
+    return ~x
+
+
+def bitwise_not(A):
+    logger.debug("GEMS_ASCEND BITWISE NOT")
+    if A.dtype == torch.bfloat16:
+        return bitwise_not_func(A.view(torch.int16)).view(torch.bfloat16)
+    return bitwise_not_func(A)
+
+
+def bitwise_not_(A):
+    logger.debug("GEMS_ASCEND BITWISE NOT_")
+    if A.dtype == torch.bfloat16:
+        bitwise_not_func(A.view(torch.int16), out0=A.view(torch.int16))
+        return A
+    bitwise_not_func(A, out0=A)
+    return A

src/flag_gems/runtime/backend/_ascend/ops/bitwise_or.py

@@ -0,0 +1,63 @@
+import logging
+
+import torch
+import triton
+
+from flag_gems.utils import pointwise_dynamic
+
+logger = logging.getLogger(f'flag_gems.runtime._ascend.ops.{__name__.split(".")[-1]}')
+
+
+@pointwise_dynamic(promotion_methods=[(0, 1, "DEFAULT")])
+@triton.jit
+def bitwise_or_func(x, y):
+    return x | y
+
+
+def bitwise_or_tensor(A, B):
+    logger.debug("GEMS_ASCEND BITWISE OR")
+    if A.dtype == torch.bfloat16:
+        return bitwise_or_func(
+            A.view(torch.int16), B.view(torch.int16)
+        ).view(torch.bfloat16)
+    return bitwise_or_func(A, B)
+
+
+def bitwise_or_tensor_(A, B):
+    logger.debug("GEMS_ASCEND BITWISE OR_")
+    if A.dtype == torch.bfloat16:
+        bitwise_or_func(A.view(torch.int16), B.view(torch.int16), out0=A.view(torch.int16))
+        return A
+    return bitwise_or_func(A, B, out0=A)
+
+
+@pointwise_dynamic(is_tensor=[True, False], promotion_methods=[(0, 1, "DEFAULT")])
+@triton.jit
+def bitwise_or_func_scalar(x, y):
+    return x | y
+
+
+def bitwise_or_scalar(A, B):
+    logger.debug("GEMS_ASCEND BITWISE OR SCALAR")
+    if A.dtype == torch.bfloat16:
+        return bitwise_or_func_scalar(
+            A.view(torch.int16), int(B)
+        ).view(torch.bfloat16)
+    return bitwise_or_func_scalar(A, B)
+
+
+def bitwise_or_scalar_(A, B):
+    logger.debug("GEMS_ASCEND BITWISE OR_ SCALAR")
+    if A.dtype == torch.bfloat16:
+        bitwise_or_func_scalar(A.view(torch.int16), int(B), out0=A.view(torch.int16))
+        return A
+    return bitwise_or_func_scalar(A, B, out0=A)
+
+
+def bitwise_or_scalar_tensor(A, B):
+    logger.debug("GEMS_ASCEND BITWISE OR SCALAR TENSOR")
+    if B.dtype == torch.bfloat16:
+        return bitwise_or_func_scalar(
+            B.view(torch.int16), int(A)
+        ).view(torch.bfloat16)
+    return bitwise_or_func_scalar(B, A)