[inductor] Fix a use before def error (#956)

Summary: for #918

Author: desertfire

test/test_torchinductor.py

@@ -2871,7 +2871,7 @@ def fn(x):
             ],
         )
 
-    def test_tmp_not_defined_issue(self):
+    def test_tmp_not_defined_issue1(self):
         def forward(
             primals_3,
             primals_4,
@@ -2920,6 +2920,22 @@ def forward(
         inps = [torch.randn(shape, dtype=dtype) for (shape, dtype) in inps]
         self.common(forward, inps)
 
+    def test_tmp_not_defined_issue2(self):
+        def forward(arg38_1, arg81_1, getitem_17, new_zeros_default_4):
+            div_tensor_7 = torch.ops.aten.div.Tensor(getitem_17, arg81_1)
+            mul_tensor_24 = torch.ops.aten.mul.Tensor(div_tensor_7, arg38_1)
+            sum_default_7 = torch.ops.aten.sum.default(mul_tensor_24)
+            return (new_zeros_default_4, sum_default_7)
+
+        args = [
+            ((1, 88, 40, 40), (140800, 1600, 40, 1), torch.float32),
+            ((), (), torch.float32),
+            ((1, 88, 40, 40), (140800, 1600, 40, 1), torch.float32),
+            ((3,), (1,), torch.float32),
+        ]
+        args = [rand_strided(shape, stride, dtype) for shape, stride, dtype in args]
+        self.common(forward, args)
+
 
 if HAS_CPU:
 

torchinductor/codegen/triton.py

@@ -616,7 +616,12 @@ def combine_contiguous_dims(self, index: sympy.Expr, tree: IterationRangesRoot):
         new_index = index.subs(dict(zip(index_vars, reindex(new_index_vars))))
         return new_index
 
-    def indexing(self, index: sympy.Expr, copy_shape=None, dense_indexing=False):
+    def indexing(
+        self,
+        index: sympy.Expr,
+        copy_shape=None,
+        dense_indexing=False,
+    ):
         """
         Compute the index and mask to pass to tl.load() or tl.store()
         """
@@ -632,6 +637,7 @@ def indexing(self, index: sympy.Expr, copy_shape=None, dense_indexing=False):
             or indirect_indexing
             or self._load_mask is not None
         ) and index != 0
+
         have_dense = True
         have_loop_vars = False
         mask = []
@@ -646,7 +652,7 @@ def indexing(self, index: sympy.Expr, copy_shape=None, dense_indexing=False):
                 mask.append(f"{tree.prefix}mask")
             dense_mask.append(f"{tree.prefix}mask")
 
-        if need_dense and not have_dense:
+        if (need_dense and not have_dense) or index == 0:
             mask = dense_mask
             index_str = f"{index_str} + tl.zeros({self.dense_size_str()}, tl.int32)"
         elif not have_loop_vars and copy_shape:
@@ -704,7 +710,12 @@ def load(self, name: str, index: sympy.Expr, upcast: bool = False):
         if upcast:
             line += ".to(tl.float32)"
 
-        if self.inside_reduction and "rmask" not in mask and not indirect_indexing:
+        if (
+            self.inside_reduction
+            and "rmask" not in mask
+            and "tmp" not in mask
+            and not indirect_indexing
+        ):
             # can lift a common load outside of reduction loop
             # One exception is when this is an indirect_load.
             tmp = self.cse.generate(self.body, line)