[Iluvatar] fix kron scatter randperm quantile

src/flag_gems/ops/randperm.py

@@ -215,7 +215,7 @@ def radix_sortbykey_scatter_kernel(
             + ((portion_id * passes + p) * max_tiles_per_portion + pid0) * bins
             + bin_id,
             partial_counter,
-            cache_modifier=".wt",
+            cache_modifier=".cg",
         )
         bin_offset = p * (bins + 1) + bin_id
         prefix_offsets = tl.load(
@@ -242,7 +242,7 @@ def radix_sortbykey_scatter_kernel(
             + ((portion_id * passes + p) * max_tiles_per_portion + pid0) * bins
             + bin_id,
             global_counter,
-            cache_modifier=".wt",
+            cache_modifier=".cg",
         )
         inc_bucket_offset = prefix_offsets.to(tl.int64) + inc_sum.to(tl.int64)
         if last_block and portion_id < num_portions - 1:

src/flag_gems/ops/scatter.py

@@ -36,7 +36,7 @@ def generate_scatter_kernel(
 
     code.writeline("def heur_block(args):")
     with code.indent():
-        code.writeline("if(flag_gems.vendor_name=='metax'):")
+        code.writeline("if(flag_gems.vendor_name in ['metax', 'iluvatar']):")
         with code.indent():
             code.writeline("return 256")
         code.writeline("return 128")

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

@@ -4,6 +4,6 @@
     vendor_name="iluvatar", device_name="cuda", device_query_cmd="ixsmi"
 )
 
-CUSTOMIZED_UNUSED_OPS = ("scatter", "quantile", "randperm", "mv")
+CUSTOMIZED_UNUSED_OPS = ()
 
 __all__ = ["*"]

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

@@ -1,5 +1,10 @@
 from .bmm import bmm
-from .div import div_mode, floor_divide, remainder, true_divide
+from .div import div_mode, div_mode_
 from .mm import mm
 
-__all__ = ["bmm", "mm", "div_mode", "floor_divide", "remainder", "true_divide"]
+__all__ = [
+    "bmm",
+    "mm",
+    "div_mode",
+    "div_mode_",
+]

src/flag_gems/runtime/backend/_iluvatar/ops/div.py

@@ -43,6 +43,14 @@ def true_divide(A, B):
         return torch.tensor(A / B)
 
 
+def true_divide_(A, B):
+    logging.debug("GEMS TRUE_DIVIDE_")
+    if isinstance(B, torch.Tensor):
+        return true_div_func(A, B, out0=A)
+    else:
+        return true_div_func_tensor_scalar(A, B, out0=A)
+
+
 @pointwise_dynamic(promotion_methods=[(0, 1, "DEFAULT")])
 @triton.jit
 def trunc_div_func(x, y):
@@ -62,7 +70,7 @@ def trunc_div_func_scalar_tensor(x, y):
 
 
 def trunc_divide(A, B):
-    logging.debug("GEMS TRUNC_DIVIDE iluvatar")
+    logging.debug("GEMS TRUNC_DIVIDE")
     if isinstance(A, torch.Tensor) and isinstance(B, torch.Tensor):
         return trunc_div_func(A, B)
     elif isinstance(A, torch.Tensor):
@@ -74,6 +82,14 @@ def trunc_divide(A, B):
         return torch.tensor(A / B)
 
 
+def trunc_divide_(A, B):
+    logging.debug("GEMS TRUNC_DIVIDE_")
+    if isinstance(B, torch.Tensor):
+        return trunc_div_func(A, B, out0=A)
+    else:
+        return trunc_div_func_tensor_scalar(A, B, out0=A)
+
+
 @triton.jit
 def _int_floordiv(x, y):
     # TODO: request Triton to add an integer remainder builtin
@@ -167,6 +183,14 @@ def floor_divide(A, B):
         return torch.tensor(A // B)
 
 
+def floor_divide_(A, B):
+    logging.debug("GEMS FLOOR_DIVIDE_")
+    if isinstance(B, torch.Tensor):
+        return floor_div_func(A, B, out0=A)
+    else:
+        return floor_div_func_tensor_scalar(A, B, out0=A)
+
+
 def div_mode(A, B, rounding_mode=None):
     if rounding_mode is None:
         return true_divide(A, B)
@@ -179,6 +203,18 @@ def div_mode(A, B, rounding_mode=None):
         raise ValueError(msg)
 
 
+def div_mode_(A, B, rounding_mode=None):
+    if rounding_mode is None:
+        return true_divide_(A, B)
+    elif rounding_mode == "trunc":
+        return trunc_divide_(A, B)
+    elif rounding_mode == "floor":
+        return floor_divide_(A, B)
+    else:
+        msg = f"div expected rounding_mode to be one of None, 'trunc', or 'floor' but found {rounding_mode}."
+        raise ValueError(msg)
+
+
 @triton.jit
 def _remainder(x, y):
     r = x % y
@@ -216,3 +252,11 @@ def remainder(A, B):
     else:
         # Both scalar
         return torch.tensor(A % B)
+
+
+def remainder_(A, B):
+    logging.debug("GEMS REMAINDER_")
+    if isinstance(B, torch.Tensor):
+        return rem_tt(A, B, out0=A)
+    else:
+        return rem_ts(A, B, out0=A)

src/flag_gems/runtime/backend/_iluvatar/tune_configs.yaml

@@ -3415,3 +3415,21 @@ batch_norm:
   - 8
   - 16
   - 32
+kron:
+- gen: true
+  param_map:
+    META:
+      BLOCK_M: block_m
+      BLOCK_N: block_n
+    num_warps: warps
+  block_m:
+  - 1
+  - 2
+  - 4
+  - 8
+  block_n:
+  - 1024
+  - 2048
+  warps:
+  - 4
+  - 8

tests/test_binary_pointwise_ops.py

@@ -620,7 +620,7 @@ def test_accuracy_trunc_div_(shape, dtype):
 
     inp1 = torch.randn(shape, dtype=dtype, device="cpu").to(flag_gems.device)
     inp2 = torch.randn(shape, dtype=dtype, device="cpu").to(flag_gems.device)
-    upcast = True if flag_gems.vendor_name not in ["kunlunxin"] else False
+    upcast = True if flag_gems.vendor_name not in ["kunlunxin", "iluvatar"] else False
     ref_inp1 = to_reference(inp1, upcast)
     ref_inp2 = to_reference(inp2, upcast)