[inductor] ignore block ptr advancements for removed buffers (pytorch#148087)

Follow up to pytorch#147193. Some buffers are removed only when the kernel context is exited so defer the lines instead.

Added `use_block_ptr` as a parameter to test case that fails if run with block ptrs enabled.

Fixes #ISSUE_NUMBER

Pull Request resolved: pytorch#148087
Approved by: https://github.com/jansel, https://github.com/eellison

Author: kundaMwiza

test/inductor/test_torchinductor.py

@@ -816,10 +816,10 @@ def wrapper(self):
 
 def skip_if_not_triton(fn):
     @functools.wraps(fn)
-    def wrapper(self):
+    def wrapper(self, *args, **kwargs):
         if not is_triton_backend(self.device):
             raise unittest.SkipTest(f"triton backend is required for {self.device}")
-        return fn(self)
+        return fn(self, *args, **kwargs)
 
     return wrapper
 
@@ -9614,7 +9614,11 @@ def fn(x):
             ],
         )
 
-    def test_tmp_not_defined_issue1(self):
+    @parametrize(
+        "use_block_ptr",
+        [subtest(False), subtest(True, decorators=[skip_if_not_triton])],
+    )
+    def test_tmp_not_defined_issue1(self, use_block_ptr):
         def forward(
             primals_3,
             primals_4,
@@ -9651,7 +9655,8 @@ def forward(
             (torch.Size([1, 512, 1]), torch.float32),
         ]
         inps = [torch.randn(shape, dtype=dtype) for (shape, dtype) in inps]
-        self.common(forward, inps, atol=1e-05, rtol=2e-05)
+        with config.patch("triton.use_block_ptr", use_block_ptr):
+            self.common(forward, inps, atol=1e-05, rtol=2e-05)
 
     @unittest.skipIf(
         os.environ.get("BUILD_ENVIRONMENT", "").startswith("parallelnative"),

torch/_inductor/codegen/triton.py

@@ -74,7 +74,6 @@
     DeferredLine,
     IndentedBuffer,
     InplacedBuffer,
-    is_buffer_removed,
     OpOverrides,
     PythonPrinter,
     RemovedArg,
@@ -3142,8 +3141,6 @@ def codegen_body(self):
                     for block_ptr, advancement in self.pointer_advancements[
                         tree.symt
                     ].items():
-                        if is_buffer_removed(self.block_ptr_to_buffer[block_ptr]):
-                            continue
                         # Subtract any advancements made in the previous loop level.
                         if level < len(loop_trees) - 1:
                             prev_tree = loop_trees[level + 1]
@@ -3158,7 +3155,10 @@ def codegen_body(self):
                             ]
 
                         self.body.writeline(
-                            f"{block_ptr} = tl.advance({block_ptr}, {V.kernel.index_to_str(advancement)})"
+                            DeferredLine(
+                                self.block_ptr_to_buffer[block_ptr],
+                                f"{block_ptr} = tl.advance({block_ptr}, {V.kernel.index_to_str(advancement)})",
+                            )
                         )
 
                 # Invalidate any cache entries that came from inside the loop.