[SILOptimizer] Add ‘llvm.cttz.*’ Intrinsic support to Constant Folding

rdar://problem/29522851

Author: Joe Shajrawi

lib/SILOptimizer/Utils/ConstantFolding.cpp

@@ -278,6 +278,44 @@ constantFoldBinaryWithOverflow(BuiltinInst *BI, BuiltinValueKind ID,
            ResultsInError);
 }
 
+static SILValue countZeros(BuiltinInst *BI, llvm::Intrinsic::ID ID) {
+  assert(BI->getArguments().size() == 2 && "Ctlz should have 2 args.");
+  OperandValueArrayRef Args = BI->getArguments();
+
+  // Fold for integer constant arguments.
+  auto *LHS = dyn_cast<IntegerLiteralInst>(Args[0]);
+  if (!LHS) {
+    return nullptr;
+  }
+  APInt LHSI = LHS->getValue();
+  unsigned LZ = 0;
+  // Check corner-case of source == zero
+  if (LHSI == 0) {
+    auto *RHS = dyn_cast<IntegerLiteralInst>(Args[1]);
+    if (!RHS || RHS->getValue() != 0) {
+      // Undefined
+      return nullptr;
+    }
+    LZ = LHSI.getBitWidth();
+  } else {
+    switch (ID) {
+    default:
+      return nullptr;
+    case llvm::Intrinsic::ctlz: {
+      LZ = LHSI.countLeadingZeros();
+      break;
+    }
+    case llvm::Intrinsic::cttz: {
+      LZ = LHSI.countTrailingZeros();
+      break;
+    }
+    }
+  }
+  APInt LZAsAPInt = APInt(LHSI.getBitWidth(), LZ);
+  SILBuilderWithScope B(BI);
+  return B.createIntegerLiteral(BI->getLoc(), LHS->getType(), LZAsAPInt);
+}
+
 static SILValue constantFoldIntrinsic(BuiltinInst *BI, llvm::Intrinsic::ID ID,
                                       Optional<bool> &ResultsInError) {
   switch (ID) {
@@ -291,31 +329,9 @@ static SILValue constantFoldIntrinsic(BuiltinInst *BI, llvm::Intrinsic::ID ID,
     return Op1;
   }
 
-  case llvm::Intrinsic::ctlz: {
-    assert(BI->getArguments().size() == 2 && "Ctlz should have 2 args.");
-    OperandValueArrayRef Args = BI->getArguments();
-
-    // Fold for integer constant arguments.
-    auto *LHS = dyn_cast<IntegerLiteralInst>(Args[0]);
-    if (!LHS) {
-      return nullptr;
-    }
-    APInt LHSI = LHS->getValue();
-    unsigned LZ = 0;
-    // Check corner-case of source == zero
-    if (LHSI == 0) {
-      auto *RHS = dyn_cast<IntegerLiteralInst>(Args[1]);
-      if (!RHS || RHS->getValue() != 0) {
-        // Undefined
-        return nullptr;
-      }
-      LZ = LHSI.getBitWidth();
-    } else {
-      LZ = LHSI.countLeadingZeros();
-    }
-    APInt LZAsAPInt = APInt(LHSI.getBitWidth(), LZ);
-    SILBuilderWithScope B(BI);
-    return B.createIntegerLiteral(BI->getLoc(), LHS->getType(), LZAsAPInt);
+  case llvm::Intrinsic::ctlz:
+  case llvm::Intrinsic::cttz: {
+    return countZeros(BI, ID);
   }
 
   case llvm::Intrinsic::sadd_with_overflow:

test/SILOptimizer/constant_propagation.sil

@@ -50,6 +50,36 @@ bb0:
 // CHECK-NEXT: return [[RES]] : $Builtin.Int64
 }
 
+sil @count_trailing_zeros_corner_case : $@convention(thin) () -> Builtin.Int64 {
+bb0:
+ %zero64 = integer_literal $Builtin.Int64, 0
+ %zero1 = integer_literal $Builtin.Int1, 0
+ %cttz = builtin "int_cttz_Int64"(%zero64 : $Builtin.Int64, %zero1 : $Builtin.Int1) : $Builtin.Int64
+ return %cttz : $Builtin.Int64
+
+// CHECK-LABEL: sil @count_trailing_zeros_corner_case
+// CHECK-NOT: integer_literal $Builtin.Int64, 0
+// CHECK-NOT: integer_literal $Builtin.Int1, 0
+// CHECK-NOT: builtin
+// CHECK: [[RES:%.*]] = integer_literal $Builtin.Int64, 64
+// CHECK-NEXT: return [[RES]] : $Builtin.Int64
+}
+
+sil @count_trailing_zeros : $@convention(thin) () -> Builtin.Int64 {
+bb0:
+ %zero64 = integer_literal $Builtin.Int64, 2
+ %zero1 = integer_literal $Builtin.Int1, 0
+ %cttz = builtin "int_cttz_Int64"(%zero64 : $Builtin.Int64, %zero1 : $Builtin.Int1) : $Builtin.Int64
+ return %cttz : $Builtin.Int64
+
+// CHECK-LABEL: sil @count_trailing_zeros
+// CHECK-NOT: integer_literal $Builtin.Int64, 2
+// CHECK-NOT: integer_literal $Builtin.Int1, 0
+// CHECK-NOT: builtin
+// CHECK: [[RES:%.*]] = integer_literal $Builtin.Int64, 1
+// CHECK-NEXT: return [[RES]] : $Builtin.Int64
+}
+
 // Compute an expression using a chain of arithmetic with overflow instructions: 2 * (2 + 3) - 3
 sil @fold_arithmetic_with_overflow : $@convention(thin) () -> Builtin.Int64 {
 bb0: