Change how index 0 load is implemented

Author: desertfire

torchinductor/codegen/triton.py

@@ -704,11 +704,11 @@ def mask_loads(self, mask):
     def load(self, name: str, index: sympy.Expr, upcast: bool = False):
         var = self.args.input(name)
         indirect_indexing = self.is_indirect_indexing(index)
-        is_index_0 = index == 0
-        index, mask = self.indexing(index)
-        if is_index_0 and "tl.zeros" not in index:
-            # Need dense_indexing when index == 0
-            index = f"{index} + tl.zeros({self.dense_size_str()}, tl.int32)"
+        if index == 0:
+            # No need to use mask when loading a single element from index 0
+            index, mask = "0", "None"
+        else:
+            index, mask = self.indexing(index)
 
         if "rmask" in mask:
             # This eviction policy heuristic is untested.