[inductor] Add abs to index_propagation (pytorch#124616)

Pull Request resolved: pytorch#124616
Approved by: https://github.com/lezcano
ghstack dependencies: pytorch#124119

Author: peterbell10

test/inductor/test_torchinductor.py

@@ -1206,6 +1206,22 @@ def repeat(x, n):
         self.assertEqual(expect, actual)
         self.assertEqual(actual, repeat(x, 3))
 
+    def test_index_propagation_abs(self):
+        def reflection_pad_left(x, n):
+            # e.g. x=[1, 2, 3], n=2 => returns [3, 2, 1, 2, 3]
+            i = torch.arange(x.shape[0] + n, device=x.device)
+            return x[(i - n).abs()]
+
+        x = torch.randn(8, device=self.device)
+        opt_fn = torch._dynamo.optimize("inductor")(reflection_pad_left)
+
+        # this should be collapsed to direct indexing
+        actual = _run_and_assert_no_indirect_indexing(
+            self, opt_fn, x, 3, has_wrapping=False
+        )
+        expect = reflection_pad_left(x, 3)
+        self.assertEqual(expect, actual)
+
     @skipIfRocm
     @config.patch(debug_index_asserts=False)
     def test_neg_index(self):

torch/_inductor/index_propagation.py

@@ -82,6 +82,10 @@ def to_dtype(
     ) -> TypedExpr:
         return TypedExpr(value.expr, dtype)
 
+    @staticmethod
+    def abs(x: TypedExpr) -> TypedExpr:
+        return TypedExpr(abs(x.expr), x.dtype)  # type: ignore[arg-type]
+
     @staticmethod
     def square(x: TypedExpr) -> TypedExpr:
         return TypedExpr(x.expr * x.expr, x.dtype)