[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 ↗

(On Diff #297832)

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

784

Nit: drop braces

804–807

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/Constants.h

public:
   static Constant *getSDiv(Constant *C1, Constant *C2, bool isExact = false);
   static Constant *getFDiv(Constant *C1, Constant *C2);
   static Constant *getURem(Constant *C1, Constant *C2);
   static Constant *getSRem(Constant *C1, Constant *C2);
   static Constant *getFRem(Constant *C1, Constant *C2);
   static Constant *getAnd(Constant *C1, Constant *C2);
   static Constant *getOr(Constant *C1, Constant *C2);
   static Constant *getXor(Constant *C1, Constant *C2);
-  static Constant *getUMin(Constant *C1, Constant *C2);
   static Constant *getShl(Constant *C1, Constant *C2,
                           bool HasNUW = false, bool HasNSW = false);
   static Constant *getLShr(Constant *C1, Constant *C2, bool isExact = false);
   static Constant *getAShr(Constant *C1, Constant *C2, bool isExact = false);
   static Constant *getTrunc(Constant *C, Type *Ty, bool OnlyIfReduced = false);
   static Constant *getSExt(Constant *C, Type *Ty, bool OnlyIfReduced = false);
   static Constant *getZExt(Constant *C, Type *Ty, bool OnlyIfReduced = false);
   static Constant *getFPTrunc(Constant *C, Type *Ty,
public:
   static Constant *getExactAShr(Constant *C1, Constant *C2) {
     return getAShr(C1, C2, true);
   }
 
   static Constant *getExactLShr(Constant *C1, Constant *C2) {
     return getLShr(C1, C2, true);
   }
 
+  /// Return a constant with unsigned clamped values between MinValue and
+  /// MaxValue inclusive: MinValue <= C <= MaxValue.
+  /// If PassThruUndefs is true then undefs from C will be returned as undef,
+  /// otherwise they will be set to MinValue.
+  static Constant *getUClamp(Constant *C, const APInt &MinValue,
+                             const APInt &MaxValue, bool PassThruUndefs = false);
+
   /// Return the identity constant for a binary opcode.
   /// The identity constant C is defined as X op C = X and C op X = X for every
   /// X when the binary operation is commutative. If the binop is not
   /// commutative, callers can acquire the operand 1 identity constant by
   /// setting AllowRHSConstant to true. For example, any shift has a zero
   /// identity constant for operand 1: X shift 0 = X.
   /// Return nullptr if the operator does not have an identity constant.
   static Constant *getBinOpIdentity(unsigned Opcode, Type *Ty,
public:
   );
 
   /// Create a FPExt, Bitcast or FPTrunc for fp -> fp casts
   static Constant *getFPCast(
     Constant *C,    ///< The integer constant to be casted
     Type *Ty ///< The integer type to cast to
   );
 
+  /// Return a bool or bool-vector set to true where the element was undef
+  /// and false otherwise. The type will be a i1 or <X x i1> matching the
+  /// number of elements in the source type. If the source type is scalable
+  /// it can only be resolved if C is an UndefValue.
+  /// May return nullptr if C is an unsupported type.
+  static Constant *getUndefMask(Constant *C);
+
   /// Return true if this is a convert constant expression
   bool isCast() const;
 
   /// Return true if this is a compare constant expression
   bool isCompare() const;
 
   /// Return true if this is an insertvalue or extractvalue expression,
   /// and the getIndices() method may be used.

llvm/lib/IR/Constants.cpp

while (I != E) {
     // If the constant was dead, then the iterator is invalidated.
     if (LastNonDeadUser == E)
       I = user_begin();
     else
       I = std::next(LastNonDeadUser);
   }
 }
 
+Constant *ConstantExpr::getUndefMask(Constant *C) {
+  assert(C && "Expected non-nullptr constant argument");
+  Type *Ty = C->getType();
+  assert((Ty->isIntOrIntVectorTy() || Ty->isFPOrFPVectorTy()) &&
+         "Only int/fp scalar/vectors supported");
+
+  Type *BoolSclTy = IntegerType::getInt1Ty(C->getContext());
+  Type *BoolTy = BoolSclTy;
+  if (Ty->isVectorTy())
+    BoolTy = VectorType::get(BoolSclTy, cast<VectorType>(Ty));
+
+  // Entire constant is undef - return true;
+  if (match(C, m_Undef()))
+    return ConstantInt::getTrue(BoolTy);
+
+  // Scalar isn't undef - return false.
+  if (!Ty->isVectorTy())
+    return ConstantInt::getFalse(BoolTy);
+
+  // Splats contains no undefs - return false;
+  if (C->getSplatValue(/*AllowUndefs*/ false))
+    return ConstantInt::getFalse(BoolTy);
+
+  // Don't know how to deal with this constant.
+  auto *VTy = dyn_cast<FixedVectorType>(Ty);
+  if (!VTy)
+    return nullptr;
+
+  unsigned NumElts = VTy->getNumElements();
+  Constant *TrueElt = ConstantInt::getTrue(BoolSclTy);
+  Constant *FalseElt = ConstantInt::getFalse(BoolSclTy);
+  SmallVector<Constant *, 32> NewC(NumElts);
+  for (unsigned i = 0; i != NumElts; ++i) {
+    Constant *EltC = C->getAggregateElement(i);
+    if (!EltC)
+      return nullptr;
+    NewC[i] = match(EltC, m_Undef()) ? TrueElt : FalseElt;
+  }
+  return ConstantVector::get(NewC);
+}
+
 Constant *Constant::replaceUndefsWith(Constant *C, Constant *Replacement) {
   assert(C && Replacement && "Expected non-nullptr constant arguments");
   Type *Ty = C->getType();
   if (match(C, m_Undef())) {
     assert(Ty == Replacement->getType() && "Expected matching types");
     return Replacement;
   }
 
for (unsigned i = 0; i != NumElts; ++i) {
     NewC[i] = EltC && match(EltC, m_Undef()) ? Replacement : EltC;
   }
   return ConstantVector::get(NewC);
 }
 
 
 //===----------------------------------------------------------------------===//
 //                                ConstantInt
 //===----------------------------------------------------------------------===//
 
 ConstantInt::ConstantInt(IntegerType *Ty, const APInt &V)
     : ConstantData(Ty, ConstantIntVal), Val(V) {
   assert(V.getBitWidth() == Ty->getBitWidth() && "Invalid constant for type");
 }
 
 ConstantInt *ConstantInt::getTrue(LLVMContext &Context) {
   LLVMContextImpl *pImpl = Context.pImpl;
   if (!pImpl->TheTrueVal)
     pImpl->TheTrueVal = ConstantInt::get(Type::getInt1Ty(Context), 1);
   return pImpl->TheTrueVal;
 }
 
 ConstantInt *ConstantInt::getFalse(LLVMContext &Context) {
   LLVMContextImpl *pImpl = Context.pImpl;
   if (!pImpl->TheFalseVal)
     pImpl->TheFalseVal = ConstantInt::get(Type::getInt1Ty(Context), 0);
   return pImpl->TheFalseVal;
 }
 
 Constant *ConstantInt::getTrue(Type *Ty) {
   assert(Ty->isIntOrIntVectorTy(1) && "Type not i1 or vector of i1.");
   ConstantInt *TrueC = ConstantInt::getTrue(Ty->getContext());
   if (auto *VTy = dyn_cast<VectorType>(Ty))
     return ConstantVector::getSplat(VTy->getElementCount(), TrueC);
   return TrueC;
 }
 
 Constant *ConstantInt::getFalse(Type *Ty) {
   assert(Ty->isIntOrIntVectorTy(1) && "Type not i1 or vector of i1.");
   ConstantInt *FalseC = ConstantInt::getFalse(Ty->getContext());

 Constant *ConstantExpr::getOr(Constant *C1, Constant *C2) {
   return get(Instruction::Or, C1, C2);
 }
 
 Constant *ConstantExpr::getXor(Constant *C1, Constant *C2) {
   return get(Instruction::Xor, C1, C2);
 }
 
-Constant *ConstantExpr::getUMin(Constant *C1, Constant *C2) {
-  Constant *Cmp = ConstantExpr::getICmp(CmpInst::ICMP_ULT, C1, C2);
-  return getSelect(Cmp, C1, C2);
-}
-
 Constant *ConstantExpr::getShl(Constant *C1, Constant *C2,
                                bool HasNUW, bool HasNSW) {
   unsigned Flags = (HasNUW ? OverflowingBinaryOperator::NoUnsignedWrap : 0) |
                    (HasNSW ? OverflowingBinaryOperator::NoSignedWrap   : 0);
   return get(Instruction::Shl, C1, C2, Flags);
 }
 
 Constant *ConstantExpr::getLShr(Constant *C1, Constant *C2, bool isExact) {
   return get(Instruction::LShr, C1, C2,
              isExact ? PossiblyExactOperator::IsExact : 0);
 }
 
 Constant *ConstantExpr::getAShr(Constant *C1, Constant *C2, bool isExact) {
   return get(Instruction::AShr, C1, C2,
              isExact ? PossiblyExactOperator::IsExact : 0);
 }
 
+Constant *ConstantExpr::getUClamp(Constant *C, const APInt &MinValue,
+                                  const APInt &MaxValue, bool PassThruUndefs) {
+  assert(MinValue.ule(MaxValue) && "Illegal clamp range");
+  Type *Ty = C->getType();
+  Constant *MinC = ConstantInt::get(Ty, MinValue);
+  Constant *MaxC = ConstantInt::get(Ty, MaxValue);
+
+  // By using UGE then UndefValue in C will be returned as true,
+  // we then select with MinC as the true value guaranteeing that they are set.
+  Constant *Result = C;
+  Constant *MinCmp = ConstantExpr::getICmp(CmpInst::ICMP_UGE, MinC, Result);
+  Result = getSelect(MinCmp, MinC, Result);
+  Constant *MaxCmp = ConstantExpr::getICmp(CmpInst::ICMP_ULE, MaxC, Result);
+  Result = getSelect(MaxCmp, MaxC, Result);
+
+  if (PassThruUndefs) {
+    Constant *IsUndef = ConstantExpr::getUndefMask(C);
+    assert(IsUndef && "Unknown undefs");
+    Result = getSelect(IsUndef, UndefValue::get(Ty), Result);
+  }
+
+  return Result;
+}
+
 Constant *ConstantExpr::getBinOpIdentity(unsigned Opcode, Type *Ty,
                                          bool AllowRHSConstant) {
   assert(Instruction::isBinaryOp(Opcode) && "Only binops allowed");
 
   // Commutative opcodes: it does not matter if AllowRHSConstant is set.
   if (Instruction::isCommutative(Opcode)) {
     switch (Opcode) {
       case Instruction::Add: // X + 0 = X

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);
+        Constant *ShAmt =
+            ConstantExpr::getUClamp(C, APInt::getNullValue(SrcWidth),
+                                    APInt(SrcWidth, DestWidth - 1), true);
         ShAmt = ConstantExpr::getTrunc(ShAmt, A->getType());
         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);
+        Constant *ShAmt =
+            ConstantExpr::getUClamp(C, APInt::getNullValue(SrcWidth),
+                                    APInt(SrcWidth, AWidth - 1), true);
         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);
+      if (auto *IsUndef = ConstantExpr::getUndefMask(C))
+        ShAmt = ConstantExpr::getSelect(IsUndef, UndefValue::get(A->getType()),
+                                        ShAmt);
       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
 }