Add D8 integer specialization tier (divisibility = 8 elements)

python/src/specialize.cc

@@ -54,6 +54,7 @@ static PyObject *u64_str = nullptr;
 static PyObject *fp32_str = nullptr;
 static PyObject *u1_str = nullptr;
 static PyObject *D_str = nullptr;
+static PyObject *D8_str = nullptr;
 static PyObject *constexpr_str = nullptr;
 static PyObject *empty_str = nullptr;
 static PyObject *nvTmaDesc_str = nullptr;
@@ -103,6 +104,7 @@ void init_interned_strings() {
   fp32_str = intern_from_string("fp32");
   u1_str = intern_from_string("u1");
   D_str = intern_from_string("D");
+  D8_str = intern_from_string("D8");
   constexpr_str = intern_from_string("constexpr");
   empty_str = intern_from_string("");
   nvTmaDesc_str = intern_from_string("nvTmaDesc");
@@ -280,7 +282,12 @@ std::pair<py::object, py::object> handle_long_type(PyObject *backend,
   if (overflow == 0) {
     type_str = (val >= INT32_MIN && val <= INT32_MAX) ? i32_str : i64_str;
     if (specialize_value) {
-      key_obj = (align && ((val & 15) == 0)) ? D_str : empty_str;
+      if (align && ((val & 15) == 0))
+        key_obj = D_str;
+      else if (align && ((val & 7) == 0))
+        key_obj = D8_str;
+      else
+        key_obj = empty_str;
     }
   } else {
     unsigned long long val_64 = PyLong_AsUnsignedLongLong(arg);
@@ -296,7 +303,12 @@ std::pair<py::object, py::object> handle_long_type(PyObject *backend,
     }
     type_str = u64_str;
     if (specialize_value) {
-      key_obj = (align && ((val_64 & 15) == 0)) ? D_str : empty_str;
+      if (align && ((val_64 & 15) == 0))
+        key_obj = D_str;
+      else if (align && ((val_64 & 7) == 0))
+        key_obj = D8_str;
+      else
+        key_obj = empty_str;
     }
   }
   if (!key_obj) {

python/test/unit/cuda/test_tensor_descriptor_cuda.py

@@ -18,3 +18,16 @@ def kernel(a, b):
     h = kernel.warmup(desc, 16, grid=(1, ))
     assert "%a: !tt.tensordesc<128xf32>" in h.asm["ttir"]
     assert "%b: i32 {tt.divisibility = 16 : i32}" in h.asm["ttir"]
+
+
+def test_int_arg_d8_specialization():
+
+    @triton.jit
+    def kernel(a, b, c):
+        pass
+
+    # a=16 -> tt.divisibility = 16, b=72 -> tt.divisibility = 8, c=3 -> no hint.
+    h = kernel.warmup(16, 72, 3, grid=(1, ))
+    assert "%a: i32 {tt.divisibility = 16 : i32}" in h.asm["ttir"]
+    assert "%b: i32 {tt.divisibility = 8 : i32}" in h.asm["ttir"]
+    assert "%c: i32 {tt.divisibility" not in h.asm["ttir"]

python/test/unit/runtime/test_specialize.py

@@ -92,17 +92,25 @@ def native_inputs_to_specialize():
         None,
         False,
         True,
-        1,
-        0,
+        -2 * 31 - 1,
         -1,
+        0,
+        1,
+        8,
         16,
         17,
+        24,
+        56,
+        72,
+        2**31 - 8,
         2**31 - 1,
         2**31,
-        -2 * 31 - 1,
+        2**63 - 8,
         2**63 - 1,
         2**63,
         2**63 + 1,
+        2**63 + 8,
+        2**64 - 8,
         2**64 - 1,
     ]
 
@@ -175,3 +183,37 @@ def test_specialize_impl(input_generator, backend, is_const, specialize_value, a
         result = native_specialize_impl(backend, arg, is_const, specialize_value, align)
         expected = reference_specialize_impl(backend, arg, is_const, specialize_value, align)
         assert result == expected
+
+
+@pytest.mark.parametrize("backend", [CUDABackend, HIPBackend])
+@pytest.mark.parametrize("value,expected_key", [
+    # Divisible by 16 -> the standard `D` specialization (divisibility = 16).
+    (0, "D"),
+    (16, "D"),
+    (-32, "D"),
+    (2**31, "D"),
+    # Divisible by 8 but not 16 -> the `D8` tier (divisibility = 8).
+    (8, "D8"),
+    (24, "D8"),
+    (72, "D8"),
+    (2**31 - 8, "D8"),
+    # Not divisible by 8 -> no specialization key.
+    (3, ""),
+    (7, ""),
+    (17, ""),
+])
+def test_d8_int_specialization(backend, value, expected_key):
+    """`D` maps to divisibility = 16, `D8` to divisibility = 8."""
+    assert backend.get_int_specialization(value, align=True) == expected_key
+    # No alignment hint requested -> the function must not specialize.
+    assert backend.get_int_specialization(value, align=False) == ""
+    # The string descriptor round-trips through `parse_attr` into the right
+    # `tt.divisibility` attribute value (or no attribute for the empty key).
+    attrs = backend.parse_attr(expected_key)
+    divisibility = next((v for k, v in attrs if k == "tt.divisibility"), None)
+    if expected_key == "D":
+        assert divisibility == 16
+    elif expected_key == "D8":
+        assert divisibility == 8
+    else:
+        assert divisibility is None

python/triton/backends/compiler.py

@@ -74,15 +74,25 @@ def get_module_map(self) -> Dict[str, ModuleType]:
     @staticmethod
     def parse_attr(desc):
         assert isinstance(desc, str)
+        # Descriptors are concatenations of single-letter flags ("D"/"D8" for
+        # divisibility, plus backend-specific suffixes like AMD's "S" for
+        # pointer_range). Check the longer "D8" prefix first so the "D" check
+        # doesn't also match "D8".
         ret = []
-        if "D" in desc:
+        if "D8" in desc:
+            ret += [["tt.divisibility", 8]]
+        elif "D" in desc:
             ret += [["tt.divisibility", 16]]
         return ret
 
     @staticmethod
     def get_int_specialization(arg, **kwargs):
-        if arg % 16 == 0 and kwargs.get("align", False):
+        if not kwargs.get("align", False):
+            return ""
+        if arg % 16 == 0:
             return "D"
+        if arg % 8 == 0:
+            return "D8"
         return ""
 
     @staticmethod