[BACKEND] Add iluvatar compute spec key specialization

python/triton/runtime/jit.py

@@ -413,7 +413,9 @@ def create_function_from_signature(sig, kparams):
     }
 
     func_namespace['mangle_type'] = mangle_type
-    func_namespace['compute_spec_key'] = compute_spec_key
+
+    from triton.runtime.driver import spec_func
+    func_namespace['compute_spec_key'] = spec_func("compute_spec_key") or compute_spec_key
 
     # Execute the function string in func_namespace to create the function
     exec(func_body, func_namespace)

third_party/iluvatar/backend/spec/__init__.py

@@ -61,4 +61,5 @@
     "bmm",
     "language_modify_all",
     "corex_sme",
+    "compute_spec_key",
 ]

third_party/iluvatar/backend/spec/triton/runtime/jit.py

@@ -56,6 +56,17 @@ def ext_JITFunction_spec_of(arg):
     return (arg % 4 == 0, arg % JITFunction.divisibility_8 == 0, arg == 1)
 
 
+def compute_spec_key(v):
+    if hasattr(v, "data_ptr") and (v.data_ptr() % 4 == 0):
+        return "D"
+    elif isinstance(v, int):
+        if v % 4 == 0:
+            return "D"
+        elif v == 1:
+            return "1"
+    return "N"
+
+
 def is_corex_param(x, enable_sme):
     if enable_sme:
         if hasattr(x, "data_ptr"):