put index expression in parens, to work around triton bug (#1504)

* put index expression in parens, to work around triton bug

* test

test/test_torchinductor.py

@@ -2918,6 +2918,15 @@ def fn(a, dim, index, b):
             rtol=1e-3,
         )
 
+    def test_scatter4(self):
+        def fn(x, ind, src):
+            return torch.scatter(x, 0, ind, src)
+
+        self.common(
+            fn,
+            (torch.randn(196, 992), torch.randint(196, (1, 992)), torch.randn(1, 992)),
+        )
+
     @unittest.skip("Flaky test, needs debugging")
     def test_scatter_add1(self):
         def fn(a, dim, index, b):

torchinductor/codegen/triton.py

@@ -772,7 +772,7 @@ def load(self, name: str, index: sympy.Expr):
             other = ", other=0"
         else:
             other = ""
-        line = f"tl.load({var} + {index}, {mask}{ep}{other})"
+        line = f"tl.load({var} + ({index}), {mask}{ep}{other})"
         if V.graph.get_dtype(name) in (torch.float16, torch.bfloat16):
             line += ".to(tl.float32)"
 
@@ -796,9 +796,9 @@ def store(self, name, index, value, mode=None):
         var = self.args.output(name)
         index, mask = self.indexing(index, value, dense_indexing=True)
         if mode is None:
-            line = f"tl.store({var} + {index}, {value}, {mask})"
+            line = f"tl.store({var} + ({index}), {value}, {mask})"
         elif mode == "atomic_add":
-            line = f"tl.atomic_add({var} + {index}, {value}, {mask})"
+            line = f"tl.atomic_add({var} + ({index}), {value}, {mask})"
         else:
             raise NotImplementedError(f"store mode={mode}")
         self.stores.writeline(name, line)