[X86][InstCombine] Add some constant folding for BZHI intrinsics

[craig.topper]

Summary

This intrinsic clears the upper bits starting at a specified index. If the index is a constant we can do some simplifications.

This could be in InstSimplify, but we don't handle any target specific intrinsics there today.

[zvi]

LGTM with minor comments.

lib/Transforms/InstCombine/InstCombineCalls.cpp

2281 ↗

(On Diff #108849)

Minor: this is not constant-folding - it's a simplification, so consider updating the comment.

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

212 ↗

(On Diff #108849)

This is out of scope, but would combining to 'and i32 %a, 0xefffffff' be profitable? It would at least be friendlier to DemandedBits.

[RKSimon]

lib/Transforms/InstCombine/InstCombineCalls.cpp

2282 ↗

(On Diff #108849)

Can we support this case? It guarantees that the CF flag is set.

[RKSimon]

LGTM

lib/Transforms/InstCombine/InstCombineCalls.cpp

2282 ↗

(On Diff #108849)

Sorry that was rubbish - yes we can fold this.

llvm/trunk/lib/Transforms/InstCombine/InstCombineCalls.cpp

if (auto *C = dyn_cast<ConstantInt>(II->getArgOperand(1))) {
         Result &= maskTrailingOnes<uint64_t>(Length);
         return replaceInstUsesWith(CI, ConstantInt::get(II->getType(), Result));
       }
       // TODO should we turn this into 'and' if shift is 0? Or 'shl' if we
       // are only masking bits that a shift already cleared?
     }
     break;
 
+  case Intrinsic::x86_bmi_bzhi_32:
+  case Intrinsic::x86_bmi_bzhi_64:
+    // If the RHS is a constant we can try some simplifications.
+    if (auto *C = dyn_cast<ConstantInt>(II->getArgOperand(1))) {
+      uint64_t Index = C->getZExtValue() & 0xff;
+      unsigned BitWidth = II->getType()->getIntegerBitWidth();
+      if (Index >= BitWidth)
+        return replaceInstUsesWith(CI, II->getArgOperand(0));
+      if (Index == 0)
+        return replaceInstUsesWith(CI, ConstantInt::get(II->getType(), 0));
+      // If the LHS is also a constant, we can completely constant fold this.
+      if (auto *InC = dyn_cast<ConstantInt>(II->getArgOperand(0))) {
+        uint64_t Result = InC->getZExtValue();
+        Result &= maskTrailingOnes<uint64_t>(Index);
+        return replaceInstUsesWith(CI, ConstantInt::get(II->getType(), Result));
+      }
+      // TODO should we convert this to an AND if the RHS is constant?
+    }
+    break;
+
   case Intrinsic::x86_vcvtph2ps_128:
   case Intrinsic::x86_vcvtph2ps_256: {
     auto Arg = II->getArgOperand(0);
     auto ArgType = cast<VectorType>(Arg->getType());
     auto RetType = cast<VectorType>(II->getType());
     unsigned ArgWidth = ArgType->getNumElements();
     unsigned RetWidth = RetType->getNumElements();
     assert(RetWidth <= ArgWidth && "Unexpected input/return vector widths");

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

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 
 declare i32 @llvm.x86.tbm.bextri.u32(i32, i32) nounwind readnone
 declare i64 @llvm.x86.tbm.bextri.u64(i64, i64) nounwind readnone
 declare i32 @llvm.x86.bmi.bextr.32(i32, i32) nounwind readnone
 declare i64 @llvm.x86.bmi.bextr.64(i64, i64) nounwind readnone
+declare i32 @llvm.x86.bmi.bzhi.32(i32, i32) nounwind readnone
+declare i64 @llvm.x86.bmi.bzhi.64(i64, i64) nounwind readnone
 
 define i32 @test_x86_tbm_bextri_u32(i32 %a) nounwind readnone {
 ; CHECK-LABEL: @test_x86_tbm_bextri_u32(
 ; CHECK-NEXT:    [[TMP1:%.*]] = tail call i32 @llvm.x86.tbm.bextri.u32(i32 [[A:%.*]], i32 1296)
 ; CHECK-NEXT:    ret i32 [[TMP1]]
 ;
   %1 = tail call i32 @llvm.x86.tbm.bextri.u32(i32 %a, i32 1296)
   ret i32 %1

 
 define i64 @test_x86_bmi_bextri_64_constfold3() nounwind readnone {
 ; CHECK-LABEL: @test_x86_bmi_bextri_64_constfold3(
 ; CHECK-NEXT:    ret i64 233495534
 ;
   %1 = tail call i64 @llvm.x86.bmi.bextr.64(i64 3735928559, i64 32772) ; extract bits 131:4 from 0xDEADBEEF
   ret i64 %1
 }
+
+define i32 @test_x86_bmi_bzhi_32(i32 %a) nounwind readnone {
+; CHECK-LABEL: @test_x86_bmi_bzhi_32(
+; CHECK-NEXT:    [[TMP1:%.*]] = tail call i32 @llvm.x86.bmi.bzhi.32(i32 [[A:%.*]], i32 31)
+; CHECK-NEXT:    ret i32 [[TMP1]]
+;
+  %1 = tail call i32 @llvm.x86.bmi.bzhi.32(i32 %a, i32 31)
+  ret i32 %1
+}
+
+define i32 @test_x86_bmi_bzhi_32_zero(i32 %a) nounwind readnone {
+; CHECK-LABEL: @test_x86_bmi_bzhi_32_zero(
+; CHECK-NEXT:    ret i32 0
+;
+  %1 = tail call i32 @llvm.x86.bmi.bzhi.32(i32 %a, i32 0)
+  ret i32 %1
+}
+
+define i32 @test_x86_bmi_bzhi_32_max(i32 %a) nounwind readnone {
+; CHECK-LABEL: @test_x86_bmi_bzhi_32_max(
+; CHECK-NEXT:    ret i32 [[A:%.*]]
+;
+  %1 = tail call i32 @llvm.x86.bmi.bzhi.32(i32 %a, i32 32)
+  ret i32 %1
+}
+
+define i32 @test_x86_bmi_bzhi_32_constfold() nounwind readnone {
+; CHECK-LABEL: @test_x86_bmi_bzhi_32_constfold(
+; CHECK-NEXT:    ret i32 1
+;
+  %1 = tail call i32 @llvm.x86.bmi.bzhi.32(i32 5, i32 1)
+  ret i32 %1
+}
+
+define i64 @test_x86_bmi_bzhi_64(i64 %a) nounwind readnone {
+; CHECK-LABEL: @test_x86_bmi_bzhi_64(
+; CHECK-NEXT:    [[TMP1:%.*]] = tail call i64 @llvm.x86.bmi.bzhi.64(i64 [[A:%.*]], i64 63)
+; CHECK-NEXT:    ret i64 [[TMP1]]
+;
+  %1 = tail call i64 @llvm.x86.bmi.bzhi.64(i64 %a, i64 63)
+  ret i64 %1
+}
+
+define i64 @test_x86_bmi_bzhi_64_zero(i64 %a) nounwind readnone {
+; CHECK-LABEL: @test_x86_bmi_bzhi_64_zero(
+; CHECK-NEXT:    ret i64 0
+;
+  %1 = tail call i64 @llvm.x86.bmi.bzhi.64(i64 %a, i64 0)
+  ret i64 %1
+}
+
+define i64 @test_x86_bmi_bzhi_64_max(i64 %a) nounwind readnone {
+; CHECK-LABEL: @test_x86_bmi_bzhi_64_max(
+; CHECK-NEXT:    ret i64 [[A:%.*]]
+;
+  %1 = tail call i64 @llvm.x86.bmi.bzhi.64(i64 %a, i64 64)
+  ret i64 %1
+}
+
+define i64 @test_x86_bmi_bzhi_64_constfold() nounwind readnone {
+; CHECK-LABEL: @test_x86_bmi_bzhi_64_constfold(
+; CHECK-NEXT:    ret i64 1
+;
+  %1 = tail call i64 @llvm.x86.bmi.bzhi.64(i64 5, i64 1)
+  ret i64 %1
+}