[llvm] 7510f32 - [MachineSink] Fix crash due to use-after-free in a MachineInstr* cache.

[Amara Emerson via llvm-commits]

Author: Amara Emerson
Date: 2023-10-06T15:02:39-07:00
New Revision: 7510f32f906ab4e583542eae2611b020f88629af

URL: https://github.com/llvm/llvm-project/commit/7510f32f906ab4e583542eae2611b020f88629af
DIFF: https://github.com/llvm/llvm-project/commit/7510f32f906ab4e583542eae2611b020f88629af.diff

LOG: [MachineSink] Fix crash due to use-after-free in a MachineInstr* cache.

After the SinkAndFold optimization was enabled, we saw some crashes with
GISel due to SinkAndFold erasing an MI while a reference was being held in a
cache.

Added: 
    llvm/test/CodeGen/AArch64/machine-sink-cache-invalidation.ll

Modified: 
    llvm/lib/CodeGen/MachineSink.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/CodeGen/MachineSink.cpp b/llvm/lib/CodeGen/MachineSink.cpp
index 2d9ff33d2755a15..40f33d664414dbb 100644
--- a/llvm/lib/CodeGen/MachineSink.cpp
+++ b/llvm/lib/CodeGen/MachineSink.cpp
@@ -540,6 +540,8 @@ bool MachineSinking::PerformSinkAndFold(MachineInstr &MI,
     }
     LLVM_DEBUG(dbgs() << "yielding"; New->dump());
     SinkDst->eraseFromParent();
+    // Clear the StoreInstrCache, since we may have invalidated it by erasing.
+    StoreInstrCache.clear();
   }
 
   // Collect operands that need to be cleaned up because the registers no longer

diff  --git a/llvm/test/CodeGen/AArch64/machine-sink-cache-invalidation.ll b/llvm/test/CodeGen/AArch64/machine-sink-cache-invalidation.ll
new file mode 100644
index 000000000000000..ce000021fb29bff
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/machine-sink-cache-invalidation.ll
@@ -0,0 +1,51 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 3
+; RUN: llc < %s -mtriple=aarch64 -global-isel | FileCheck %s
+
+target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
+target triple = "arm64-apple-macosx13.5.0"
+
+; Check we don't crash here with a use-after-free.
+
+define i32 @nsis_BZ2_bzDecompress(ptr %pos.i, i1 %cmp661.not3117.i, i1 %exitcond.not.i) {
+; CHECK-LABEL: nsis_BZ2_bzDecompress:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    b .LBB0_2
+; CHECK-NEXT:  .LBB0_1: // %while.end671.i
+; CHECK-NEXT:    // in Loop: Header=BB0_2 Depth=1
+; CHECK-NEXT:    strb w8, [x0]
+; CHECK-NEXT:    tbnz w2, #0, .LBB0_4
+; CHECK-NEXT:  .LBB0_2: // %for.body653.i
+; CHECK-NEXT:    // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    ldrb w8, [x0]
+; CHECK-NEXT:    tbnz w1, #0, .LBB0_1
+; CHECK-NEXT:  // %bb.3: // %while.body663.i
+; CHECK-NEXT:    // in Loop: Header=BB0_2 Depth=1
+; CHECK-NEXT:    mov x9, xzr
+; CHECK-NEXT:    ldrb w9, [x9]
+; CHECK-NEXT:    strb wzr, [x0, x9]
+; CHECK-NEXT:    b .LBB0_1
+; CHECK-NEXT:  .LBB0_4: // %for.end677.i
+; CHECK-NEXT:    mov w0, wzr
+; CHECK-NEXT:    ret
+entry:
+  br label %for.body653.i
+
+for.body653.i:                                    ; preds = %while.end671.i, %entry
+  %0 = load i8, ptr null, align 1
+  %idxprom657.i2 = zext i8 %0 to i64
+  %1 = load i8, ptr %pos.i, align 1
+  br i1 %cmp661.not3117.i, label %while.end671.i, label %while.body663.i
+
+while.body663.i:                                  ; preds = %for.body653.i
+  %arrayidx669.i = getelementptr [6 x i8], ptr %pos.i, i64 0, i64 %idxprom657.i2
+  store i8 0, ptr %arrayidx669.i, align 1
+  br label %while.end671.i
+
+while.end671.i:                                   ; preds = %while.body663.i, %for.body653.i
+  store i8 %1, ptr %pos.i, align 1
+  br i1 %exitcond.not.i, label %for.end677.i, label %for.body653.i
+
+for.end677.i:                                     ; preds = %while.end671.i
+  ret i32 0
+}
+