support.tl.extra.shim.dispatch.to.backends

Author: kxyk99

src/flag_gems/ops/div.py

@@ -4,12 +4,9 @@
 import triton
 import triton.language as tl
 
-from ..utils import pointwise_dynamic, tl_extra_shim
+from ..utils import pointwise_dynamic
+from ..utils.triton_lang_extension import div_rn, div_rz, fmod, trunc
 
-div_rn = tl_extra_shim.div_rn
-div_rz = tl_extra_shim.div_rz
-fmod = tl_extra_shim.fmod
-trunc = tl_extra_shim.trunc
 logger = logging.getLogger(__name__)
 
 

src/flag_gems/ops/silu.py

@@ -3,9 +3,8 @@
 import triton
 import triton.language as tl
 
-from ..utils import pointwise_dynamic, tl_extra_shim
-
-div_rn = tl_extra_shim.div_rn
+from ..utils import pointwise_dynamic
+from ..utils.triton_lang_extension import div_rn
 
 logger = logging.getLogger(__name__)
 

src/flag_gems/utils/triton_lang_extension.py

@@ -14,6 +14,8 @@
 import triton
 from triton import language as tl
 
+from .triton_lang_helper import use_tl_extra
+
 
 @triton.jit
 def program_id(
@@ -70,3 +72,34 @@ def maximum_with_index_tie_break_right(a_value, a_index, b_value, b_index):
     # Prefer highest index if values are equal
     mask |= equal & (a_index > b_index)
     return tl.where(mask, a_value, b_value), tl.where(mask, a_index, b_index)
+
+
+@triton.jit
+@use_tl_extra
+def div_rn(x, y):
+    """div_rn default - round to nearest"""
+    result = x / y
+    return tl.floor(result + 0.5)
+
+
+@triton.jit
+@use_tl_extra
+def div_rz(x, y):
+    """div_rz default - round toward zero"""
+    result = x / y
+    return tl.where(result >= 0, tl.floor(result), tl.ceil(result))
+
+
+@triton.jit
+@use_tl_extra
+def fmod(x, y):
+    """fmod default - floating point modulo"""
+    quotient = div_rz(x, y)
+    return x - y * quotient
+
+
+@triton.jit
+@use_tl_extra
+def trunc(x):
+    """trunc default - truncate to integer"""
+    return tl.where(x >= 0, tl.floor(x), tl.ceil(x))

src/flag_gems/utils/triton_lang_helper.py

@@ -20,3 +20,23 @@
         tl_extra_shim = triton.language.math
     except ImportError:
         tl_extra_shim = triton.language.libdevice
+
+
+def use_backend(module):
+    """using backend module impl"""
+
+    def decorator(func):
+        func_name = func.__name__
+        if hasattr(module, func_name):
+            try:
+                return getattr(module, func_name)
+            except Exception:
+                pass
+        return func
+
+    return decorator
+
+
+def use_tl_extra(func):
+    """backend function shim"""
+    return use_backend(tl_extra_shim)(func)