Fix crash in SymInt unary minus (#116160)

malfet · pytorchmergebot · commit f2c1fb3ee400 · 2023-12-20T20:12:57.000Z
Before this change `-SymInt(std::numeric_limits<int64_t>::min()) == 0` would reliably crash with null pointer dereference, as `data_` of the SymInt returned by `operator-` would be `0x8000000000000000`, because of the carry/overflow flags set by `negq`. Before the change x86_64 assembly generated for https://github.com/pytorch/pytorch/blob/4f02cc06706648799a33218711e62eae859d5927/c10/core/SymInt.cpp#L137 looked as follows: ``` 0x7ffff7f2f490 <+115>: movq %rax, %rdx 0x7ffff7f2f493 <+118>: negq %rdx 0x7ffff7f2f496 <+121>: movq %rdx, (%rbp) 0x7ffff7f2f49a <+125>: movabsq $0x4000000000000000, %rdx ; imm = 0x4000000000000000 0x7ffff7f2f4a4 <+135>: cmpq %rdx, %rax 0x7ffff7f2f4a7 <+138>: jle 0x7ffff7f2f520 ; <+259> at SymInt.cpp:141:1 ``` `negq %rfx` correspond to unary minus and `cmpq %rdx, 0x4000000000000000` are inverted `check_range` https://github.com/pytorch/pytorch/blob/b6d0d0819ae6099b2d659d03892e1e19c17e912b/c10/core/SymInt.h#L247-L249 Flags raised by `negq` will affect the results of `cmpq`, and as result value would not be allocated on heap, but rather preserved as `nullptr`. Not sure if it's worth benchmarking, but perhaps using `__builtin_sub_overflow` would be faster as it does not require an extra comparison, just guarantees that overflow flags is cleared after the op. Pull Request resolved: #116160 Approved by: https://github.com/Skylion007, https://github.com/colesbury
diff --git a/c10/core/SymInt.cpp b/c10/core/SymInt.cpp
@@ -134,7 +134,13 @@ bool SymInt::expect_size(const char* file, int64_t line) const {
 
 SymInt operator-(const SymInt& s) {
   if (auto ma = s.maybe_as_int()) {
-    return SymInt(-*ma);
+    const auto val = *ma;
+    // Note: Result of `-std::numeric_limits<decltype(val)>::min()` is undefined
+    // But on many platforms it equals to self + setting Carry/Overflow flags
+    // Which in opimized code affects results of `check_range` condition
+    // Workaround by using ternary that avoids alterning the flags
+    constexpr auto val_min = std::numeric_limits<decltype(val)>::min();
+    return SymInt(val != val_min ? -val : val_min);
   } else {
     return SymInt(s.toSymNodeImplUnowned()->neg());
   }
diff --git a/c10/test/core/SymInt_test.cpp b/c10/test/core/SymInt_test.cpp
@@ -22,4 +22,13 @@ TEST(SymIntTest, CheckRange) {
   EXPECT_FALSE(SymInt::check_range(INT64_MIN));
 }
 
+TEST(SymIntTest, Overflows) {
+  const auto x = SymInt(INT64_MAX);
+  EXPECT_NE(-(x + 1), 0);
+
+  const auto y = SymInt(INT64_MIN);
+  EXPECT_NE(-y, 0);
+  EXPECT_NE(0 - y, 0);
+}
+
 #endif
