[InstCombine] visitTrunc - pass through undefs for trunc(shift(trunc/ext(x),c)) patterns

[RKSimon]

Summary

Based on the recent patches D88475 and D88429 where we losing undef values due to extension/comparisons.

I've added a Constant::mergeUndefsWith method that merges the undef scalar/elements from another Constant into a specific Constant.

[lebedev.ri]

Instead of eventually potentially duplicating the entirety of ConstantExpr methods,
i wonder if we can get away with an inverse of Constant::replaceUndefsWith(),
that takes two constants, and does an element-wise or'ing of undefs,
i.e. the LHS constant gets all the additional undefs that the RHS constant has.

[RKSimon]

Reworked with a Constant::mergeUndefsWith helper as suggested by @lebedev.ri

[lebedev.ri]

Perfect, LGTM.

llvm/include/llvm/IR/Constant.h

208

I'm having trouble parsing the first line.
I guess it talks about vector element undefs and scalar undefs, but it's hard go get that.

llvm/lib/IR/Constants.cpp

743

Nit: drop braces

763–766

You probably want

if (!match(NewC[I], m_Undef()) && match(OtherEltC, m_Undef())) {
  NewC[I] = OtherEltC;
  FoundExtraUndef = true;
}

otherwise you set FoundExtraUndef even if LHS already was undef.

llvm/include/llvm/IR/Constant.h

public:
   Constant *stripPointerCasts() {
     return const_cast<Constant*>(
                       static_cast<const Constant *>(this)->stripPointerCasts());
   }
 
   /// Try to replace undefined constant C or undefined elements in C with
   /// Replacement. If no changes are made, the constant C is returned.
   static Constant *replaceUndefsWith(Constant *C, Constant *Replacement);
+
+  /// Merges undefs of a Constant with another Constant, along with the
+  /// undefs already present. Other doesn't have to be the same type as C, but
+  /// both must either be scalars or vectors with the same element count. If no
+  /// changes are made, the constant C is returned.
+  static Constant *mergeUndefsWith(Constant *C, Constant *Other);
 };
 
 } // end namespace llvm
 
 #endif // LLVM_IR_CONSTANT_H

llvm/lib/IR/Constants.cpp

for (unsigned i = 0; i != NumElts; ++i) {
     Constant *EltC = C->getAggregateElement(i);
     assert((!EltC || EltC->getType() == Replacement->getType()) &&
            "Expected matching types");
     NewC[i] = EltC && match(EltC, m_Undef()) ? Replacement : EltC;
   }
   return ConstantVector::get(NewC);
 }
 
+Constant *Constant::mergeUndefsWith(Constant *C, Constant *Other) {
+  assert(C && Other && "Expected non-nullptr constant arguments");
+  if (match(C, m_Undef()))
+    return C;
+
+  Type *Ty = C->getType();
+  if (match(Other, m_Undef()))
+    return UndefValue::get(Ty);
+
+  auto *VTy = dyn_cast<FixedVectorType>(Ty);
+  if (!VTy)
+    return C;
+
+  Type *EltTy = VTy->getElementType();
+  unsigned NumElts = VTy->getNumElements();
+  assert(isa<FixedVectorType>(Other->getType()) &&
+         cast<FixedVectorType>(Other->getType())->getNumElements() == NumElts &&
+         "Type mismatch");
+
+  bool FoundExtraUndef = false;
+  SmallVector<Constant *, 32> NewC(NumElts);
+  for (unsigned I = 0; I != NumElts; ++I) {
+    NewC[I] = C->getAggregateElement(I);
+    Constant *OtherEltC = Other->getAggregateElement(I);
+    assert(NewC[I] && OtherEltC && "Unknown vector element");
+    if (!match(NewC[I], m_Undef()) && match(OtherEltC, m_Undef())) {
+      NewC[I] = UndefValue::get(EltTy);
+      FoundExtraUndef = true;
+    }
+  }
+  if (FoundExtraUndef)
+    return ConstantVector::get(NewC);
+  return C;
+}
 
 //===----------------------------------------------------------------------===//
 //                                ConstantInt
 //===----------------------------------------------------------------------===//
 
 ConstantInt::ConstantInt(IntegerType *Ty, const APInt &V)
     : ConstantData(Ty, ConstantIntVal), Val(V) {
   assert(V.getBitWidth() == Ty->getBitWidth() && "Invalid constant for type");

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

Instruction *InstCombinerImpl::visitTrunc(TruncInst &Trunc) {
   if (match(Src, m_LShr(m_SExt(m_Value(A)), m_Constant(C)))) {
     unsigned AWidth = A->getType()->getScalarSizeInBits();
     unsigned MaxShiftAmt = SrcWidth - std::max(DestWidth, AWidth);
     auto *OldSh = cast<Instruction>(Src);
     bool IsExact = OldSh->isExact();
 
     // If the shift is small enough, all zero bits created by the shift are
     // removed by the trunc.
-    // TODO: Support passing through undef shift amounts - these currently get
-    // clamped to MaxAmt.
     if (match(C, m_SpecificInt_ICMP(ICmpInst::ICMP_ULE,
                                     APInt(SrcWidth, MaxShiftAmt)))) {
       // trunc (lshr (sext A), C) --> ashr A, C
       if (A->getType() == DestTy) {
         Constant *MaxAmt = ConstantInt::get(SrcTy, DestWidth - 1, false);
         Constant *ShAmt = ConstantExpr::getUMin(C, MaxAmt);
         ShAmt = ConstantExpr::getTrunc(ShAmt, A->getType());
+        ShAmt = Constant::mergeUndefsWith(ShAmt, C);
         return IsExact ? BinaryOperator::CreateExactAShr(A, ShAmt)
                        : BinaryOperator::CreateAShr(A, ShAmt);
       }
       // The types are mismatched, so create a cast after shifting:
       // trunc (lshr (sext A), C) --> sext/trunc (ashr A, C)
       if (Src->hasOneUse()) {
         Constant *MaxAmt = ConstantInt::get(SrcTy, AWidth - 1, false);
         Constant *ShAmt = ConstantExpr::getUMin(C, MaxAmt);
         ShAmt = ConstantExpr::getTrunc(ShAmt, A->getType());
         Value *Shift = Builder.CreateAShr(A, ShAmt, "", IsExact);
         return CastInst::CreateIntegerCast(Shift, DestTy, true);
       }
     }
     // TODO: Mask high bits with 'and'.
   }
 
   // trunc (*shr (trunc A), C) --> trunc(*shr A, C)
   if (match(Src, m_OneUse(m_Shr(m_Trunc(m_Value(A)), m_Constant(C))))) {
     unsigned MaxShiftAmt = SrcWidth - DestWidth;
 
     // If the shift is small enough, all zero/sign bits created by the shift are
     // removed by the trunc.
-    // TODO: Support passing through undef shift amounts - these currently get
-    // zero'd by getIntegerCast.
     if (match(C, m_SpecificInt_ICMP(ICmpInst::ICMP_ULE,
                                     APInt(SrcWidth, MaxShiftAmt)))) {
       auto *OldShift = cast<Instruction>(Src);
-      auto *ShAmt = ConstantExpr::getIntegerCast(C, A->getType(), true);
       bool IsExact = OldShift->isExact();
+      auto *ShAmt = ConstantExpr::getIntegerCast(C, A->getType(), true);
+      ShAmt = Constant::mergeUndefsWith(ShAmt, C);
       Value *Shift =
           OldShift->getOpcode() == Instruction::AShr
               ? Builder.CreateAShr(A, ShAmt, OldShift->getName(), IsExact)
               : Builder.CreateLShr(A, ShAmt, OldShift->getName(), IsExact);
       return CastInst::CreateTruncOrBitCast(Shift, DestTy);
     }
   }
 

llvm/test/Transforms/InstCombine/cast.ll

;
   %B = sext <2 x i8> %A to <2 x i32>
   %C = lshr <2 x i32> %B, <i32 6, i32 6>
   %D = trunc <2 x i32> %C to <2 x i8>
   ret <2 x i8> %D
 }
 
 define <2 x i8> @trunc_lshr_sext_uniform_undef(<2 x i8> %A) {
 ; ALL-LABEL: @trunc_lshr_sext_uniform_undef(
-; ALL-NEXT:    [[D:%.*]] = ashr <2 x i8> [[A:%.*]], <i8 6, i8 7>
+; ALL-NEXT:    [[D:%.*]] = ashr <2 x i8> [[A:%.*]], <i8 6, i8 undef>
 ; ALL-NEXT:    ret <2 x i8> [[D]]
 ;
   %B = sext <2 x i8> %A to <2 x i32>
   %C = lshr <2 x i32> %B, <i32 6, i32 undef>
   %D = trunc <2 x i32> %C to <2 x i8>
   ret <2 x i8> %D
 }
 
 define <2 x i8> @trunc_lshr_sext_nonuniform(<2 x i8> %A) {
 ; ALL-LABEL: @trunc_lshr_sext_nonuniform(
 ; ALL-NEXT:    [[D:%.*]] = ashr <2 x i8> [[A:%.*]], <i8 6, i8 2>
 ; ALL-NEXT:    ret <2 x i8> [[D]]
 ;
   %B = sext <2 x i8> %A to <2 x i32>
   %C = lshr <2 x i32> %B, <i32 6, i32 2>
   %D = trunc <2 x i32> %C to <2 x i8>
   ret <2 x i8> %D
 }
 
 define <3 x i8> @trunc_lshr_sext_nonuniform_undef(<3 x i8> %A) {
 ; ALL-LABEL: @trunc_lshr_sext_nonuniform_undef(
-; ALL-NEXT:    [[D:%.*]] = ashr <3 x i8> [[A:%.*]], <i8 6, i8 2, i8 7>
+; ALL-NEXT:    [[D:%.*]] = ashr <3 x i8> [[A:%.*]], <i8 6, i8 2, i8 undef>
 ; ALL-NEXT:    ret <3 x i8> [[D]]
 ;
   %B = sext <3 x i8> %A to <3 x i32>
   %C = lshr <3 x i32> %B, <i32 6, i32 2, i32 undef>
   %D = trunc <3 x i32> %C to <3 x i8>
   ret <3 x i8> %D
 }
 

llvm/test/Transforms/InstCombine/trunc-shift-trunc.ll

;
   %b = trunc <2 x i64> %a to <2 x i32>
   %c = lshr <2 x i32> %b, <i32 8, i32 2>
   %d = trunc <2 x i32> %c to <2 x i8>
   ret <2 x i8> %d
 }
 
 define <2 x i8> @trunc_lshr_trunc_uniform_undef(<2 x i64> %a) {
 ; CHECK-LABEL: @trunc_lshr_trunc_uniform_undef(
-; CHECK-NEXT:    [[C1:%.*]] = lshr <2 x i64> [[A:%.*]], <i64 24, i64 0>
+; CHECK-NEXT:    [[C1:%.*]] = lshr <2 x i64> [[A:%.*]], <i64 24, i64 undef>
 ; CHECK-NEXT:    [[D:%.*]] = trunc <2 x i64> [[C1]] to <2 x i8>
 ; CHECK-NEXT:    ret <2 x i8> [[D]]
 ;
   %b = trunc <2 x i64> %a to <2 x i32>
   %c = lshr <2 x i32> %b, <i32 24, i32 undef>
   %d = trunc <2 x i32> %c to <2 x i8>
   ret <2 x i8> %d
 }
;
   %b = trunc <2 x i64> %a to <2 x i32>
   %c = ashr <2 x i32> %b, <i32 0, i32 23>
   %d = trunc <2 x i32> %c to <2 x i8>
   ret <2 x i8> %d
 }
 
 define <2 x i8> @trunc_ashr_trunc_uniform_undef(<2 x i64> %a) {
 ; CHECK-LABEL: @trunc_ashr_trunc_uniform_undef(
-; CHECK-NEXT:    [[C1:%.*]] = ashr <2 x i64> [[A:%.*]], <i64 8, i64 0>
+; CHECK-NEXT:    [[C1:%.*]] = ashr <2 x i64> [[A:%.*]], <i64 8, i64 undef>
 ; CHECK-NEXT:    [[D:%.*]] = trunc <2 x i64> [[C1]] to <2 x i8>
 ; CHECK-NEXT:    ret <2 x i8> [[D]]
 ;
   %b = trunc <2 x i64> %a to <2 x i32>
   %c = ashr <2 x i32> %b, <i32 8, i32 undef>
   %d = trunc <2 x i32> %c to <2 x i8>
   ret <2 x i8> %d
 }