[inductor][3/N] triton support post-#5512, tt.divisibility format (#145575)

1. Fix the tt.divisibility format in hints.py. Previously, it was `{((0,), (1,)): [["tt.divisibility", 16]]}`. Now it is `{(0,): [["tt.divisibility", 16]], (1,): [["tt.divisibility", 16]]}`. This was an oversight in the first PR I added. I've verified that we now get `{ tt.divisibility = 16 }` in the generated TTGIR.
2. Update the test_codegen_triton.py test to work with multiple triton versions (and test this divisibility format in the new triton version)

Pull Request resolved: https://github.com/pytorch/pytorch/pull/145575
Approved by: https://github.com/SamGinzburg

test/inductor/test_codegen_triton.py

@@ -33,49 +33,67 @@ class TestCodegenTriton(InductorTestCase):
 
     @inductor_config.patch("triton.divisible_by_16", True)
     def test_config_of_sizearg(self):
+        from torch._inductor.utils import (
+            get_triton_attrs_descriptor_version,
+            TritonAttrsDescriptorVersion,
+        )
+
         two = sympy.Integer(2)
         eight = sympy.Integer(8)
         sixteen = sympy.Integer(16)
         s0 = sympy.Symbol("s0", positive=True, integer=True)
         s1 = sympy.Symbol("s1", positive=True, integer=True)
 
-        def _check_divisibility(config):
-            try:
-                from triton.backends.compiler import AttrsDescriptor  # noqa: F401
-
-                return config.divisibility_16
-            except ImportError:
-                return config.divisible_by_16
-
-        self.assertEqual(
-            (2,),
-            _check_divisibility(
-                triton_utils.config_of(
-                    [
-                        SizeArg("A", two),  # no
-                        SizeArg("B", eight),  # no
-                        SizeArg("C", sixteen),  # yes
-                        SizeArg("D", s0),  # no
-                        SizeArg("E", s1),  # no
-                    ]
+        def _check_divisibility(expected_divisible_indices, config):
+            if get_triton_attrs_descriptor_version() in {
+                TritonAttrsDescriptorVersion.V1_COMPILER,
+                TritonAttrsDescriptorVersion.V0_NO_TRITON,
+            }:
+                self.assertEqual(expected_divisible_indices, config.divisible_by_16)
+            elif get_triton_attrs_descriptor_version() in {
+                TritonAttrsDescriptorVersion.V2_BACKENDS,
+                TritonAttrsDescriptorVersion.V3_BACKENDS_TUPLE,
+            }:
+                self.assertEqual(expected_divisible_indices, config.divisibility_16)
+            else:
+                assert (
+                    get_triton_attrs_descriptor_version()
+                    == TritonAttrsDescriptorVersion.V4_DICT
                 )
+                self.assertIsInstance(config, dict)
+
+                for idx in expected_divisible_indices:
+                    # config is in the form
+                    # {(idx,): [["tt.divisibility", 16]]}
+                    # where (idx,) is a tuple in order to support tuple inputs to triton kernels.
+                    self.assertTrue((idx,) in config)
+                    self.assertTrue(["tt.divisibility", 16] in config[(idx,)])
+
+        _check_divisibility(
+            (2,),
+            triton_utils.config_of(
+                [
+                    SizeArg("A", two),  # no
+                    SizeArg("B", eight),  # no
+                    SizeArg("C", sixteen),  # yes
+                    SizeArg("D", s0),  # no
+                    SizeArg("E", s1),  # no
+                ]
             ),
         )
 
-        self.assertEqual(
+        _check_divisibility(
             (0, 2, 4, 5, 6),
-            _check_divisibility(
-                triton_utils.config_of(
-                    [
-                        SizeArg("A", two * eight),  # 0: yes
-                        SizeArg("B", eight * s0),  # 1: no
-                        SizeArg("C", two * eight * s0),  # 2: yes
-                        SizeArg("D", s0 * s1),  # 3: no
-                        SizeArg("E", sixteen * s0),  # 4: yes
-                        SizeArg("F", sixteen * eight * s0 * s1),  # 5: yes
-                        SizeArg("G", two * eight * s0 * s1),  # 6: yes
-                    ]
-                )
+            triton_utils.config_of(
+                [
+                    SizeArg("A", two * eight),  # 0: yes
+                    SizeArg("B", eight * s0),  # 1: no
+                    SizeArg("C", two * eight * s0),  # 2: yes
+                    SizeArg("D", s0 * s1),  # 3: no
+                    SizeArg("E", sixteen * s0),  # 4: yes
+                    SizeArg("F", sixteen * eight * s0 * s1),  # 5: yes
+                    SizeArg("G", two * eight * s0 * s1),  # 6: yes
+                ]
             ),
         )
 

torch/_inductor/runtime/hints.py

@@ -95,9 +95,7 @@ if _is_triton_available():
             divisible_by_16=None,
             equal_to_1=None,
         ):
-            return {
-                tuple((x,) for x in divisible_by_16): [["tt.divisibility", 16]],
-            }
+            return {(x,): [["tt.divisibility", 16]] for x in divisible_by_16}
 
 else:
     # Define a namedtuple as a fallback when AttrsDescriptor is not available