[LoopUnswitch] Properly update MSSA if header has non-clobbering stores.

This patch fixes updating MemorySSA if the header contains memory
defs that do not clobber a duplicated instruction. We need to find the
first defining access outside the loop body and use that as defining
access of the duplicated instruction.

This fixes a crash caused by bee486851c.

llvm/lib/Transforms/Scalar/LoopUnswitch.cpp

@@ -1114,12 +1114,16 @@ void LoopUnswitch::emitPreheaderBranchOnCondition(
 
         Loop *L = LI->getLoopFor(I->getParent());
         auto *DefiningAccess = MemA->getDefiningAccess();
-        // If the defining access is a MemoryPhi in the header, get the incoming
-        // value for the pre-header as defining access.
-        if (DefiningAccess->getBlock() == I->getParent()) {
+        // Get the first defining access before the loop.
+        while (L->contains(DefiningAccess->getBlock())) {
+          // If the defining access is a MemoryPhi, get the incoming
+          // value for the pre-header as defining access.
           if (auto *MemPhi = dyn_cast<MemoryPhi>(DefiningAccess)) {
             DefiningAccess =
                 MemPhi->getIncomingValueForBlock(L->getLoopPreheader());
+          } else {
+            DefiningAccess =
+                cast<MemoryDef>(DefiningAccess)->getDefiningAccess();
           }
         }
         MSSAU->createMemoryAccessInBB(New, DefiningAccess, New->getParent(),

llvm/test/Transforms/LoopUnswitch/partial-unswitch-update-memoryssa.ll

@@ -0,0 +1,76 @@
+; RUN: opt -loop-unswitch -verify-dom-info -verify-memoryssa -S -enable-new-pm=0 %s | FileCheck %s
+; RUN: opt -loop-unswitch -memssa-check-limit=3 -verify-dom-info -verify-memoryssa -S -enable-new-pm=0 %s | FileCheck %s
+
+declare void @clobber()
+
+; Check that MemorySSA updating can deal with a clobbering access of a
+; duplicated load being a MemoryPHI outside the loop.
+define void @partial_unswitch_memssa_update(i32* noalias %ptr, i1 %c) {
+; CHECK-LABEL: @partial_unswitch_memssa_update(
+; CHECK-LABEL: loop.ph:
+; CHECK-NEXT:    [[LV:%[a-z0-9]+]] = load i32, i32* %ptr, align 4
+; CHECK-NEXT:    [[C:%[a-z0-9]+]] = icmp eq i32 [[LV]], 0
+; CHECK-NEXT:    br i1 [[C]]
+entry:
+  br i1 %c, label %loop.ph, label %outside.clobber
+
+outside.clobber:
+  call void @clobber()
+  br label %loop.ph
+
+loop.ph:
+  br label %loop.header
+
+loop.header:
+  %lv = load i32, i32* %ptr, align 4
+  %hc = icmp eq i32 %lv, 0
+  br i1 %hc, label %if, label %then
+
+if:
+  br label %loop.latch
+
+then:
+  br label %loop.latch
+
+loop.latch:
+  br i1 true, label %loop.header, label %exit
+
+exit:
+  ret void
+}
+
+; Check that MemorySSA updating can deal with skipping defining accesses in the
+; loop body until it finds the first defining access outside the loop.
+define void @partial_unswitch_inloop_stores_beteween_outside_defining_access(i64* noalias %ptr, i16* noalias %src) {
+; CHECK-LABEL: @partial_unswitch_inloop_stores_beteween_outside_defining_access
+; CHECK-LABEL: entry:
+; CHECK-NEXT:    store i64 0, i64* %ptr, align 1
+; CHECK-NEXT:    store i64 1, i64* %ptr, align 1
+; CHECK-NEXT:    [[LV:%[a-z0-9]+]] = load i16, i16* %src, align 1
+; CHECK-NEXT:    [[C:%[a-z0-9]+]] = icmp eq i16 [[LV]], 0
+; CHECK-NEXT:    br i1 [[C]]
+;
+entry:
+  store i64 0, i64* %ptr, align 1
+  store i64 1, i64* %ptr, align 1
+  br label %loop
+
+loop:
+  %iv = phi i32 [ 0, %entry ], [ %iv.next, %loop.latch ]
+  store i64 2, i64* %ptr, align 1
+  %lv = load i16, i16* %src, align 1
+  %invar.cond = icmp eq i16 %lv, 0
+  br i1 %invar.cond, label %noclobber, label %loop.latch
+
+noclobber:
+  br label %loop.latch
+
+loop.latch:
+  %iv.next = add i32 %iv, 1
+  %ec = icmp eq i32 %iv, 1000
+  br i1 %ec, label %exit, label %loop
+
+exit:
+  ret void
+}
+

llvm/test/Transforms/LoopUnswitch/partial-unswitch.ll

@@ -575,42 +575,6 @@ exit:
   ret i32 10
 }
 
-; Check that MemorySSA updating can deal with a clobbering access of a
-; duplicated load being a MemoryPHI outside the loop.
-define void @partial_unswitch_memssa_update(i32* noalias %ptr, i1 %c) {
-; CHECK-LABEL: @partial_unswitch_memssa_update(
-; CHECK-LABEL: loop.ph:
-; CHECK-NEXT:    [[LV:%[a-z0-9]+]] = load i32, i32* %ptr, align 4
-; CHECK-NEXT:    [[C:%[a-z0-9]+]] = icmp eq i32 [[LV]], 0
-; CHECK-NEXT:    br i1 [[C]]
-entry:
-  br i1 %c, label %loop.ph, label %outside.clobber
-
-outside.clobber:
-  call void @clobber()
-  br label %loop.ph
-
-loop.ph:
-  br label %loop.header
-
-loop.header:
-  %lv = load i32, i32* %ptr, align 4
-  %hc = icmp eq i32 %lv, 0
-  br i1 %hc, label %if, label %then
-
-if:
-  br label %loop.latch
-
-then:
-  br label %loop.latch
-
-loop.latch:
-  br i1 true, label %loop.header, label %exit
-
-exit:
-  ret void
-}
-
 ; Make sure the duplicated instructions are moved to a preheader that always
 ; executes when the loop body also executes. Do not check the unswitched code,
 ; because it is already checked in the @partial_unswitch_true_successor test