[ValueTracking] peek through shuffles in ComputeNumSignBits (PR37549)

[spatel]

Summary

The motivating case is from PR37549:
https://bugs.llvm.org/show_bug.cgi?id=37549

The analysis improvement allows us to form a vector 'select' out of bitwise logic (the use of ComputeNumSignBits was added at rL345149).

The smaller test shows another InstCombine improvement - we use ComputeNumSignBits to add 'nsw' to shift-left.

[RKSimon]

lib/Analysis/ValueTracking.cpp

2515 ↗

(On Diff #170928)

Do we know for certain that we don't reference an under operand at this point? SelectionDAG::ComputeNumSignBits bails i such cases.

[spatel]

lib/Analysis/ValueTracking.cpp

2515 ↗

(On Diff #170928)

I saw that, but I wasn't sure why bailing was necessary (in IR at least).

If the shuffle's mask contains an undef, then the result for that lane is undef:
http://llvm.org/docs/LangRef.html#shufflevector-instruction

So we can choose 0 or -1 for that undef element, and so it can't reduce the number of sign bits?

[efriedma]

lib/Analysis/ValueTracking.cpp

2515 ↗

(On Diff #170928)

We can pick any value for an undef, but every optimization has to handle it consistently. So you can't assume it's 0 unless you actually replace it with zero.

In this particular case, instead of looking for UndefValues, could you examine the mask instead?

[spatel]

lib/Analysis/ValueTracking.cpp

2515 ↗

(On Diff #170928)

IIUC, there are 2 independent factors here:

1. We need to examine the mask for undef values and bail out if it contains any.
2. If we look at both vector inputs to the shuffle, then we must take the minimum signbits value of those 2 values (we already do that in the DAG version). A potential complication for this enhancement is that I'm not sure if we handle an undef input here. That's why I made that case a TODO in this patch.

[spatel]

Patch updated:
Don't try to analyze a shuffle with undef mask elements.
The negative test for 'nsw' actually demonstrates the bug in the last draft. If we had applied 'nsw' in that case, some other analysis could see poison for the undef lane that didn't exist in the original code.

[RKSimon]

LGTM with one minor

lib/Analysis/ValueTracking.cpp

2520 ↗

(On Diff #170969)

assert that both ops aren't undef?

llvm/trunk/lib/Analysis/ValueTracking.cpp

case Instruction::Trunc:
     break;
 
   case Instruction::ExtractElement:
     // Look through extract element. At the moment we keep this simple and skip
     // tracking the specific element. But at least we might find information
     // valid for all elements of the vector (for example if vector is sign
     // extended, shifted, etc).
     return ComputeNumSignBits(U->getOperand(0), Depth + 1, Q);
+
+  case Instruction::ShuffleVector:
+    // If the shuffle mask contains any undefined elements, that element of the
+    // result is undefined. Propagating information from a source operand may
+    // not be correct in that case, so just bail out.
+    if (cast<ShuffleVectorInst>(U)->getMask()->containsUndefElement())
+      break;
+
+    assert((!isa<UndefValue>(U->getOperand(0)) ||
+            !isa<UndefValue>(U->getOperand(1)))
+           && "Should have simplified shuffle with 2 undef inputs");
+
+    // Look through shuffle of 1 source vector.
+    if (isa<UndefValue>(U->getOperand(0)))
+      return ComputeNumSignBits(U->getOperand(1), Depth + 1, Q);
+    if (isa<UndefValue>(U->getOperand(1)))
+      return ComputeNumSignBits(U->getOperand(0), Depth + 1, Q);
+
+    // TODO: We can look through shuffles of 2 sources by computing the minimum
+    // sign bits for each operand (similar to what we do for binops).
+    break;
   }
 
   // Finally, if we can prove that the top bits of the result are 0's or 1's,
   // use this information.
 
   // If we can examine all elements of a vector constant successfully, we're
   // done (we can't do any better than that). If not, keep trying.
   if (unsigned VecSignBits = computeNumSignBitsVectorConstant(V, TyBits))

llvm/trunk/test/Transforms/InstCombine/logical-select.ll

 ; CHECK-NEXT:    [[CMP:%.*]] = fcmp ole <4 x float> [[X:%.*]], [[Y:%.*]]
 ; CHECK-NEXT:    [[SEXT:%.*]] = sext <4 x i1> [[CMP]] to <4 x i32>
 ; CHECK-NEXT:    [[S1:%.*]] = shufflevector <4 x i32> [[SEXT]], <4 x i32> undef, <4 x i32> <i32 0, i32 0, i32 1, i32 1>
 ; CHECK-NEXT:    [[S2:%.*]] = shufflevector <4 x i32> [[SEXT]], <4 x i32> undef, <4 x i32> <i32 2, i32 2, i32 3, i32 3>
 ; CHECK-NEXT:    [[SHUF_OR1:%.*]] = or <4 x i32> [[S1]], [[S2]]
 ; CHECK-NEXT:    [[S3:%.*]] = shufflevector <4 x i32> [[SHUF_OR1]], <4 x i32> undef, <4 x i32> <i32 0, i32 0, i32 1, i32 1>
 ; CHECK-NEXT:    [[S4:%.*]] = shufflevector <4 x i32> [[SHUF_OR1]], <4 x i32> undef, <4 x i32> <i32 2, i32 2, i32 3, i32 3>
 ; CHECK-NEXT:    [[SHUF_OR2:%.*]] = or <4 x i32> [[S3]], [[S4]]
-; CHECK-NEXT:    [[NOT_OR2:%.*]] = xor <4 x i32> [[SHUF_OR2]], <i32 -1, i32 -1, i32 -1, i32 -1>
-; CHECK-NEXT:    [[XBC:%.*]] = bitcast <4 x float> [[X]] to <4 x i32>
-; CHECK-NEXT:    [[ZBC:%.*]] = bitcast <4 x float> [[Z:%.*]] to <4 x i32>
-; CHECK-NEXT:    [[AND1:%.*]] = and <4 x i32> [[NOT_OR2]], [[XBC]]
-; CHECK-NEXT:    [[AND2:%.*]] = and <4 x i32> [[SHUF_OR2]], [[ZBC]]
-; CHECK-NEXT:    [[SEL:%.*]] = or <4 x i32> [[AND1]], [[AND2]]
-; CHECK-NEXT:    ret <4 x i32> [[SEL]]
+; CHECK-NEXT:    [[TMP1:%.*]] = trunc <4 x i32> [[SHUF_OR2]] to <4 x i1>
+; CHECK-NEXT:    [[DOTV:%.*]] = select <4 x i1> [[TMP1]], <4 x float> [[Z:%.*]], <4 x float> [[X]]
+; CHECK-NEXT:    [[TMP2:%.*]] = bitcast <4 x float> [[DOTV]] to <4 x i32>
+; CHECK-NEXT:    ret <4 x i32> [[TMP2]]
 ;
   %cmp = fcmp ole <4 x float> %x, %y
   %sext = sext <4 x i1> %cmp to <4 x i32>
   %s1 = shufflevector <4 x i32> %sext, <4 x i32> undef, <4 x i32> <i32 0, i32 0, i32 1, i32 1>
   %s2 = shufflevector <4 x i32> %sext, <4 x i32> undef, <4 x i32> <i32 2, i32 2, i32 3, i32 3>
   %shuf_or1 = or <4 x i32> %s1, %s2
   %s3 = shufflevector <4 x i32> %shuf_or1, <4 x i32> undef, <4 x i32> <i32 0, i32 0, i32 1, i32 1>
   %s4 = shufflevector <4 x i32> %shuf_or1, <4 x i32> undef, <4 x i32> <i32 2, i32 2, i32 3, i32 3>

llvm/trunk/test/Transforms/InstCombine/nsw.ll

 ; CHECK-NEXT:    ret i8 [[ADD]]
 ;
   %x = add nsw i8 %A, 10
   %y = add nsw i8 %B, 10
   %add = add nsw i8 %x, %y
   ret i8 %add
 }
 
-; TODO: ComputeNumSignBits()/computeKnownBits() should look through a shufflevector.
+; TODO: computeKnownBits() should look through a shufflevector.
 
 define <3 x i32> @shl_nuw_nsw_shuffle_splat_vec(<2 x i8> %x) {
 ; CHECK-LABEL: @shl_nuw_nsw_shuffle_splat_vec(
 ; CHECK-NEXT:    [[T2:%.*]] = zext <2 x i8> [[X:%.*]] to <2 x i32>
 ; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <2 x i32> [[T2]], <2 x i32> undef, <3 x i32> <i32 1, i32 0, i32 1>
-; CHECK-NEXT:    [[T3:%.*]] = shl <3 x i32> [[SHUF]], <i32 17, i32 17, i32 17>
+; CHECK-NEXT:    [[T3:%.*]] = shl nsw <3 x i32> [[SHUF]], <i32 17, i32 17, i32 17>
 ; CHECK-NEXT:    ret <3 x i32> [[T3]]
 ;
   %t2 = zext <2 x i8> %x to <2 x i32>
   %shuf = shufflevector <2 x i32> %t2, <2 x i32> undef, <3 x i32> <i32 1, i32 0, i32 1>
   %t3 = shl <3 x i32> %shuf, <i32 17, i32 17, i32 17>
   ret <3 x i32> %t3
 }
 
+; Negative test - if the shuffle mask contains an undef, we bail out to
+; avoid propagating information that may not be used consistently by callers.
+
+define <3 x i32> @shl_nuw_nsw_shuffle_undef_elt_splat_vec(<2 x i8> %x) {
+; CHECK-LABEL: @shl_nuw_nsw_shuffle_undef_elt_splat_vec(
+; CHECK-NEXT:    [[T2:%.*]] = zext <2 x i8> [[X:%.*]] to <2 x i32>
+; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <2 x i32> [[T2]], <2 x i32> undef, <3 x i32> <i32 1, i32 undef, i32 0>
+; CHECK-NEXT:    [[T3:%.*]] = shl <3 x i32> [[SHUF]], <i32 17, i32 17, i32 17>
+; CHECK-NEXT:    ret <3 x i32> [[T3]]
+;
+  %t2 = zext <2 x i8> %x to <2 x i32>
+  %shuf = shufflevector <2 x i32> %t2, <2 x i32> undef, <3 x i32> <i32 1, i32 undef, i32 0>
+  %t3 = shl <3 x i32> %shuf, <i32 17, i32 17, i32 17>
+  ret <3 x i32> %t3
+}
+