[instcombine][x86] Converted pdep/pext with shifted mask to simple arithmetic

[reames]

Summary

If the mask of a pdep or pext instruction is a shift masked (i.e. one contiguous block of ones) we need at most one and and one shift to represent the operation without the intrinsic. One all platforms I know of, this is faster than the pdep/pext.

The cost modelling for multiple contiguous blocks might be worth exploring in a follow up, but it's not relevant for my current use case. It would almost certainly be a win on AMDs where these are really really slow though.

[craig.topper]

LGTM

[anna]

thanks for this Philip!

llvm/lib/Target/X86/X86InstCombineIntrinsic.cpp

case Intrinsic::x86_bmi_pext_64:
     if (auto *MaskC = dyn_cast<ConstantInt>(II.getArgOperand(1))) {
       if (MaskC->isNullValue()) {
         return IC.replaceInstUsesWith(II, ConstantInt::get(II.getType(), 0));
       }
       if (MaskC->isAllOnesValue()) {
         return IC.replaceInstUsesWith(II, II.getArgOperand(0));
       }
 
+      if (MaskC->getValue().isShiftedMask()) {
+        // any single contingous sequence of 1s anywhere in the mask simply
+        // describes a subset of the input bits shifted to the appropriate
+        // position.  Replace with the straight forward IR.
+        unsigned ShiftAmount = MaskC->getValue().countTrailingZeros();
+        Value *Input = II.getArgOperand(0);
+        Value *Masked = IC.Builder.CreateAnd(Input, II.getArgOperand(1));
+        Value *Shifted = IC.Builder.CreateLShr(Masked,
+                                               ConstantInt::get(II.getType(),
+                                                                ShiftAmount));
+        return IC.replaceInstUsesWith(II, Shifted);
+      }
+
+
       if (auto *SrcC = dyn_cast<ConstantInt>(II.getArgOperand(0))) {
         uint64_t Src = SrcC->getZExtValue();
         uint64_t Mask = MaskC->getZExtValue();
         uint64_t Result = 0;
         uint64_t BitToSet = 1;
 
         while (Mask) {
           // Isolate lowest set bit.
X86TTIImpl::instCombineIntrinsic(InstCombiner &IC, IntrinsicInst &II) const {
   case Intrinsic::x86_bmi_pdep_64:
     if (auto *MaskC = dyn_cast<ConstantInt>(II.getArgOperand(1))) {
       if (MaskC->isNullValue()) {
         return IC.replaceInstUsesWith(II, ConstantInt::get(II.getType(), 0));
       }
       if (MaskC->isAllOnesValue()) {
         return IC.replaceInstUsesWith(II, II.getArgOperand(0));
       }
+      if (MaskC->getValue().isShiftedMask()) {
+        // any single contingous sequence of 1s anywhere in the mask simply
+        // describes a subset of the input bits shifted to the appropriate
+        // position.  Replace with the straight forward IR.
+        unsigned ShiftAmount = MaskC->getValue().countTrailingZeros();
+        Value *Input = II.getArgOperand(0);
+        Value *Shifted = IC.Builder.CreateShl(Input,
+                                              ConstantInt::get(II.getType(),
+                                                               ShiftAmount));
+        Value *Masked = IC.Builder.CreateAnd(Shifted, II.getArgOperand(1));
+        return IC.replaceInstUsesWith(II, Masked);
+      }
 
       if (auto *SrcC = dyn_cast<ConstantInt>(II.getArgOperand(0))) {
         uint64_t Src = SrcC->getZExtValue();
         uint64_t Mask = MaskC->getZExtValue();
         uint64_t Result = 0;
         uint64_t BitToTest = 1;
 
         while (Mask) {

llvm/test/Transforms/InstCombine/X86/x86-bmi-tbm.ll

 define i64 @test_x86_pext_64_allones_mask(i64 %x) nounwind readnone {
 ; CHECK-LABEL: @test_x86_pext_64_allones_mask(
 ; CHECK-NEXT:    ret i64 %x
 ;
   %1 = tail call i64 @llvm.x86.bmi.pext.64(i64 %x, i64 -1)
   ret i64 %1
 }
 
+define i32 @test_x86_pext_32_shifted_mask(i32 %x) nounwind readnone {
+; CHECK-LABEL: @test_x86_pext_32_shifted_mask(
+; CHECK-NEXT:    %1 = lshr i32 %x, 1
+; CHECK-NEXT:    %2 = and i32 %1, 3
+; CHECK-NEXT:    ret i32 %2
+;
+  %1 = tail call i32 @llvm.x86.bmi.pext.32(i32 %x, i32 6)
+  ret i32 %1
+}
+
+define i64 @test_x86_pext_64_shifted_mask(i64 %x) nounwind readnone {
+; CHECK-LABEL: @test_x86_pext_64_shifted_mask(
+; CHECK-NEXT:    %1 = lshr i64 %x, 1
+; CHECK-NEXT:    %2 = and i64 %1, 3
+; CHECK-NEXT:    ret i64 %2
+;
+  %1 = tail call i64 @llvm.x86.bmi.pext.64(i64 %x, i64 6)
+  ret i64 %1
+}
+
+
 define i32 @test_x86_pext_32_constant_fold() nounwind readnone {
 ; CHECK-LABEL: @test_x86_pext_32_constant_fold(
 ; CHECK-NEXT:    ret i32 30001
 ;
   %1 = tail call i32 @llvm.x86.bmi.pext.32(i32 1985229328, i32 4042322160)
   ret i32 %1
 }
 

 define i64 @test_x86_pdep_64_allones_mask(i64 %x) nounwind readnone {
 ; CHECK-LABEL: @test_x86_pdep_64_allones_mask(
 ; CHECK-NEXT:    ret i64 %x
 ;
   %1 = tail call i64 @llvm.x86.bmi.pdep.64(i64 %x, i64 -1)
   ret i64 %1
 }
 
+define i32 @test_x86_pdep_32_shifted_mask(i32 %x) nounwind readnone {
+; CHECK-LABEL: @test_x86_pdep_32_shifted_mask(
+; CHECK-NEXT:    %1 = shl i32 %x, 2
+; CHECK-NEXT:    %2 = and i32 %1, 12
+; CHECK-NEXT:    ret i32 %2
+;
+  %1 = tail call i32 @llvm.x86.bmi.pdep.32(i32 %x, i32 12)
+  ret i32 %1
+}
+
+define i64 @test_x86_pdep_64_shifted_mask(i64 %x) nounwind readnone {
+; CHECK-LABEL: @test_x86_pdep_64_shifted_mask(
+; CHECK-NEXT:    %1 = shl i64 %x, 2
+; CHECK-NEXT:    %2 = and i64 %1, 12
+; CHECK-NEXT:    ret i64 %2
+;
+  %1 = tail call i64 @llvm.x86.bmi.pdep.64(i64 %x, i64 12)
+  ret i64 %1
+}
+
+
 define i32 @test_x86_pdep_32_constant_fold() nounwind readnone {
 ; CHECK-LABEL: @test_x86_pdep_32_constant_fold(
 ; CHECK-NEXT:    ret i32 807407616
 ;
   %1 = tail call i32 @llvm.x86.bmi.pdep.32(i32 1985229328, i32 4042322160)
   ret i32 %1
 }