[InstCombine] Properly change GEP type when reassociating loop invariant GEP chains

[dneilson]

Summary

This is a fix to PR37005.

Essentially, rL328539 ([InstCombine] reassociate loop invariant GEP chains to enable LICM) contains a bug
whereby it will convert:
%src = getelementptr inbounds i8, i8* %base, <2 x i64> %val
%res = getelementptr inbounds i8, <2 x i8*> %src, i64 %val2
into:
%src = getelementptr inbounds i8, i8* %base, i64 %val2
%res = getelementptr inbounds i8, <2 x i8*> %src, <2 x i64> %val

By swapping the index operands if the GEPs are in a loop, and %val is loop variant while %val2
is loop invariant.

This fix recreates new GEP instructions if the index operand swap would result in the type
of %src changing from vector to scalar, or vice versa.

[dneilson]

lib/Transforms/InstCombine/InstructionCombining.cpp

1710

I need to change this. We should be putting the new src instruction at the same location as the one it's replacing. Otherwise, it might be possible to get into unknown weirdness.

[dneilson]

• Change insertion point of replacement src GEP instruction.

[sebpop]

Looks good. Thanks!

lib/Transforms/InstCombine/InstructionCombining.cpp

1709

Let's move SO0 before SO0Ty and use it in that computation.

[dneilson]

• Make review copy match what will be landed.

lib/Transforms/InstCombine/InstructionCombining.cpp

if (LI && Src->getNumOperands() == 2 && GEP.getNumOperands() == 2 &&
         Src->hasOneUse()) {
       if (Loop *L = LI->getLoopFor(GEP.getParent())) {
         Value *GO1 = GEP.getOperand(1);
         Value *SO1 = Src->getOperand(1);
         // Reassociate the two GEPs if SO1 is variant in the loop and GO1 is
         // invariant: this breaks the dependence between GEPs and allows LICM
         // to hoist the invariant part out of the loop.
         if (L->isLoopInvariant(GO1) && !L->isLoopInvariant(SO1)) {
+          // We have to be careful here.
+          // We have something like:
+          //  %src = getelementptr <ty>, <ty>* %base, <ty> %idx
+          //  %res = getelementptr <ty>, <ty>* %src, <ty> %idx2
+          // If we just swap idx & idx2 then we could inadvertantly
+          // change %src from a vector to a scalar, or vice versa.
+          // Cases:
+          //  1) %base a scalar & idx a scalar & idx2 a vector
+          //      => Swapping idx & idx2 turns %src into a vector type.
+          //  2) %base a scalar & idx a vector & idx2 a scalar
+          //      => Swapping idx & idx2 turns %src in a scalar type
+          //  3) %base, %idx, and %idx2 are scalars
+          //      => %src & %res are scalars
+          //      => swapping idx & idx2 is safe
+          //  4) %base a vector
+          //      => %src is a vector
+          //      => swapping idx & idx2 is safe.
+          auto *SO0Ty = Src->getOperand(0)->getType();
+          if (!isa<VectorType>(GEPType) || // case 3
+              isa<VectorType>(SO0Ty)) {    // case 4
             Src->setOperand(1, GO1);
             GEP.setOperand(1, SO1);
             return &GEP;
+          } else {
+            // Case 1 or 2
+            // -- have to recreate %src & %res
+            auto *SO0 = Src->getOperand(0);
+            Builder.SetInsertPoint(&GEP);
+            auto *NewSrc = cast<GetElementPtrInst>(
+                Builder.CreateGEP(SO0, GO1, Src->getName()));
+            NewSrc->setIsInBounds(Src->isInBounds());
+            auto *NewGEP = GetElementPtrInst::Create(nullptr, NewSrc, {SO1});
+            NewGEP->setIsInBounds(GEP.isInBounds());
+            return NewGEP;
+          }
         }
       }
     }
 
     // Note that if our source is a gep chain itself then we wait for that
     // chain to be resolved before we perform this transformation.  This
     // avoids us creating a TON of code in some cases.
     if (auto *SrcGEP = dyn_cast<GEPOperator>(Src->getOperand(0)))

test/Transforms/InstCombine/gep-combine-loop-invariant.ll

land.lhs.true:                                    ; preds = %if.then
   %dec = add i32 %chain_length.addr.08, -1
   %cmp5 = icmp eq i32 %dec, 0
   br i1 %cmp5, label %do.end, label %do.body
 
 do.end:                                           ; preds = %do.body, %land.lhs.true, %if.then, %entry
   %cont.0 = phi i32 [ 1, %entry ], [ 0, %if.then ], [ 0, %land.lhs.true ], [ 1, %do.body ]
   ret i32 %cont.0
 }
+
+declare void @blackhole(<2 x i8*>)
+
+define void @PR37005(i8* %base, i8** %in) {
+; CHECK-LABEL: @PR37005(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[E2:%.*]] = getelementptr inbounds i8*, i8** [[IN:%.*]], i64 undef
+; CHECK-NEXT:    [[E4:%.*]] = getelementptr inbounds i8*, i8** [[E2]], <2 x i64> <i64 0, i64 1>
+; CHECK-NEXT:    [[PI1:%.*]] = ptrtoint <2 x i8**> [[E4]] to <2 x i64>
+; CHECK-NEXT:    [[LR1:%.*]] = lshr <2 x i64> [[PI1]], <i64 21, i64 21>
+; CHECK-NEXT:    [[SL1:%.*]] = shl nuw nsw <2 x i64> [[LR1]], <i64 7, i64 7>
+; CHECK-NEXT:    [[E51:%.*]] = getelementptr inbounds i8, i8* [[BASE:%.*]], i64 80
+; CHECK-NEXT:    [[E6:%.*]] = getelementptr inbounds i8, i8* [[E51]], <2 x i64> [[SL1]]
+; CHECK-NEXT:    call void @blackhole(<2 x i8*> [[E6]])
+; CHECK-NEXT:    br label [[LOOP]]
+;
+entry:
+  br label %loop
+
+loop:
+  %e1 = getelementptr inbounds i8*, i8** %in, i64 undef
+  %e2 = getelementptr inbounds i8*, i8** %e1, i64 6
+  %bc1 = bitcast i8** %e2 to <2 x i8*>*
+  %e3 = getelementptr inbounds <2 x i8*>, <2 x i8*>* %bc1, i64 0, i64 0
+  %e4 = getelementptr inbounds i8*, i8** %e3, <2 x i64> <i64 0, i64 1>
+  %pi1 = ptrtoint <2 x i8**> %e4 to <2 x i64>
+  %lr1 = lshr <2 x i64> %pi1, <i64 21, i64 21>
+  %sl1 = shl nuw nsw <2 x i64> %lr1, <i64 7, i64 7>
+  %e5 = getelementptr inbounds i8, i8* %base, <2 x i64> %sl1
+  %e6 = getelementptr inbounds i8, <2 x i8*> %e5, i64 80
+  call void @blackhole(<2 x i8*> %e6)
+  br label %loop
+}
+
+define void @PR37005_2(i8* %base, i8** %in) {
+; CHECK-LABEL: @PR37005_2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[E2:%.*]] = getelementptr inbounds i8*, i8** [[IN:%.*]], i64 undef
+; CHECK-NEXT:    [[PI1:%.*]] = ptrtoint i8** [[E2]] to i64
+; CHECK-NEXT:    [[LR1:%.*]] = lshr i64 [[PI1]], 21
+; CHECK-NEXT:    [[SL1:%.*]] = shl nuw nsw i64 [[LR1]], 7
+; CHECK-NEXT:    [[E51:%.*]] = getelementptr inbounds i8, i8* [[BASE:%.*]], <2 x i64> <i64 80, i64 60>
+; CHECK-NEXT:    [[E6:%.*]] = getelementptr inbounds i8, <2 x i8*> [[E51]], i64 [[SL1]]
+; CHECK-NEXT:    call void @blackhole(<2 x i8*> [[E6]])
+; CHECK-NEXT:    br label [[LOOP]]
+;
+entry:
+  br label %loop
+
+loop:
+  %e1 = getelementptr inbounds i8*, i8** %in, i64 undef
+  %e2 = getelementptr inbounds i8*, i8** %e1, i64 6
+  %pi1 = ptrtoint i8** %e2 to i64
+  %lr1 = lshr i64 %pi1, 21
+  %sl1 = shl nuw nsw i64 %lr1, 7
+  %e5 = getelementptr inbounds i8, i8* %base, i64 %sl1
+  %e6 = getelementptr inbounds i8, i8* %e5, <2 x i64> <i64 80, i64 60>
+  call void @blackhole(<2 x i8*> %e6)
+  br label %loop
+}
+
+define void @PR37005_3(<2 x i8*> %base, i8** %in) {
+; CHECK-LABEL: @PR37005_3(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[E2:%.*]] = getelementptr inbounds i8*, i8** [[IN:%.*]], i64 undef
+; CHECK-NEXT:    [[E4:%.*]] = getelementptr inbounds i8*, i8** [[E2]], <2 x i64> <i64 0, i64 1>
+; CHECK-NEXT:    [[PI1:%.*]] = ptrtoint <2 x i8**> [[E4]] to <2 x i64>
+; CHECK-NEXT:    [[LR1:%.*]] = lshr <2 x i64> [[PI1]], <i64 21, i64 21>
+; CHECK-NEXT:    [[SL1:%.*]] = shl nuw nsw <2 x i64> [[LR1]], <i64 7, i64 7>
+; CHECK-NEXT:    [[E5:%.*]] = getelementptr inbounds i8, <2 x i8*> [[BASE:%.*]], i64 80
+; CHECK-NEXT:    [[E6:%.*]] = getelementptr inbounds i8, <2 x i8*> [[E5]], <2 x i64> [[SL1]]
+; CHECK-NEXT:    call void @blackhole(<2 x i8*> [[E6]])
+; CHECK-NEXT:    br label [[LOOP]]
+;
+entry:
+  br label %loop
+
+loop:
+  %e1 = getelementptr inbounds i8*, i8** %in, i64 undef
+  %e2 = getelementptr inbounds i8*, i8** %e1, i64 6
+  %bc1 = bitcast i8** %e2 to <2 x i8*>*
+  %e3 = getelementptr inbounds <2 x i8*>, <2 x i8*>* %bc1, i64 0, i64 0
+  %e4 = getelementptr inbounds i8*, i8** %e3, <2 x i64> <i64 0, i64 1>
+  %pi1 = ptrtoint <2 x i8**> %e4 to <2 x i64>
+  %lr1 = lshr <2 x i64> %pi1, <i64 21, i64 21>
+  %sl1 = shl nuw nsw <2 x i64> %lr1, <i64 7, i64 7>
+  %e5 = getelementptr inbounds i8, <2 x i8*> %base, <2 x i64> %sl1
+  %e6 = getelementptr inbounds i8, <2 x i8*> %e5, i64 80
+  call void @blackhole(<2 x i8*> %e6)
+  br label %loop
+}