Fix for-loop divisibility parsing (#119859)

Pull Request resolved: #119859
Approved by: https://github.com/aakhundov
ghstack dependencies: #119834, #119835, #119836, #119838

test/dynamo/test_triton_kernels.py

@@ -1149,6 +1149,40 @@ def add_4_times_kernel(
             ["out_ptr"],
         )
 
+    @make_mutation_test
+    def test_add_for_loop2():
+        @triton.jit
+        def add_1_time_kernel(
+            in_ptr0,
+            in_ptr1,
+            out_ptr,
+            n_elements,
+            BLOCK_SIZE: "tl.constexpr",
+        ):
+            pid = tl.program_id(axis=0)
+            block_start = pid * BLOCK_SIZE
+            offsets = block_start + tl.arange(0, BLOCK_SIZE)
+            mask = offsets < n_elements
+            x = tl.load(in_ptr0 + offsets, mask=mask)
+            y = tl.load(in_ptr1 + offsets, mask=mask)
+            for i in range(0, BLOCK_SIZE):
+                i = tl.multiple_of(i, 1)
+            output = x + y
+            tl.store(out_ptr + offsets, output, mask=mask)
+
+        t = torch.randn(4)
+        return (
+            add_1_time_kernel,
+            {
+                "in_ptr0": t,
+                "in_ptr1": t,
+                "out_ptr": t,
+                "n_elements": 4,
+                "BLOCK_SIZE": 4,
+            },
+            ["out_ptr"],
+        )
+
     @make_mutation_test
     def test_add_nested_for_loop():
         @triton.jit

torch/_higher_order_ops/triton_kernel_wrap.py

@@ -197,7 +197,7 @@ def parse_ttir(ttir, kwargs):
         ?stmt: op | if | for | while | condition_stmt | label_stmt | cf_stmt
 
         if: [assign_lhs "="] "scf.if" args rest stmt* "}" "else" "{" stmt* "}" LOC -> process_if
-        for: [assign_lhs "="] "scf.for" args rest stmt* "}" LOC -> process_for
+        for: [assign_lhs "="] "scf.for" args rest stmt* "}" divisibility_annot? LOC -> process_for
         while: [assign_lhs "="] "scf.while" args rest stmt* "}" "do" "{" stmt* "}" LOC -> process_while
 
         condition_stmt: "scf.condition" "(" arg ")" args rest
@@ -209,6 +209,7 @@ def parse_ttir(ttir, kwargs):
           | [assign_lhs "="] OP_NAME [FN_NAME] args rest?  -> process_op
 
         ?rest: (":" | "{" | "\\"" | "->" | "<" | "=") /.+/ NEWLINE
+        divisibility_annot: "{" "tt.divisibility_arg1" /[^}]+/ "}"
 
         args: | "(" ")" | "("? arg ("," arg)* ")"?
 
@@ -443,7 +444,7 @@ def identify_mutated_tensors(kernel, kwargs):
     except Exception as e:
         import traceback
 
-        log.info(
+        warnings.warn(
             "Encountered an exception in identify_mutated_tensors, assuming every input is mutated"
         )
         log.debug(