Integer division expansion should handle constant divisors intelligently

[dshtilman]

Summary

When expanding integer division by a constant, replace it with a simple multiplication by a magic number, rather than generating a full division loop.

Implementation follows the DAG code in TargetLowering::BuildSDIV().
Ref: "Hacker's Delight" or "The PowerPC Compiler Writer's Guide".

[joerg]

Why is this code necessary? LLVM already implements the division-by-constant optimisation.

lib/Transforms/Utils/IntegerDivision.cpp

 //===-- IntegerDivision.cpp - Expand integer division ---------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains an implementation of 32bit and 64bit scalar integer
 // division for targets that don't have native support. It's largely derived
 // from compiler-rt's implementations of __udivsi3 and __udivmoddi4,
 // but hand-tuned for targets that prefer less control flow.
+// A more optimal code is generated for division by a constant.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Utils/IntegerDivision.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Intrinsics.h"
static Value *generateUnsignedDivisionCode(Value *Dividend, Value *Divisor,
   Q_3->addIncoming(Q_1, DoWhile);
   // ;   %q_5 = phi i32 [ %q_4, %loop-exit ], [ %retVal, %special-cases ]
   Q_5->addIncoming(Q_4, LoopExit);
   Q_5->addIncoming(RetVal, SpecialCases);
 
   return Q_5;
 }
 
+/// Generate code for signed division by a constant. Implementation follows
+/// TargetLowering::BuildSDIV that replaces division with multiplication
+/// by a magic number.
+static Value *generateSignedDivisionByConstant(Value *Dividend,
+                                               ConstantInt *Divisor,
+                                               IRBuilder<> &Builder) {
+  unsigned BitWidth = Dividend->getType()->getIntegerBitWidth();
+  ConstantInt *Shift;
+  ConstantInt *Magics;
+
+  APInt d = Divisor->getValue();
+  APInt::ms magics = d.magic();
+
+  if (BitWidth == 64) {
+    Shift = Builder.getInt64(63);
+    Magics = Builder.getInt64(magics.m.getSExtValue());
+  } else {
+    assert(BitWidth == 32 && "Unexpected bit width");
+    Shift = Builder.getInt32(31);
+    Magics = Builder.getInt32(magics.m.getSExtValue());
+  }
+
+  // Multiply the numerator (operand 0) by the magic value
+  Value *Q = Builder.CreateMul(Dividend, Magics);
+
+  // If d > 0 and m < 0, add the numerator
+  if (d.isStrictlyPositive() && magics.m.isNegative()) {
+    Q = Builder.CreateAdd(Q, Dividend);
+  }
+
+  // If d < 0 and m > 0, subtract the numerator.
+  if (d.isNegative() && magics.m.isStrictlyPositive()) {
+    Q = Builder.CreateSub(Q, Dividend);
+  }
+
+  // Shift right algebraic if shift value is nonzero
+  if (magics.s > 0) {
+    ConstantInt *MagicsShift = (BitWidth == 64) ?
+      Builder.getInt64(magics.s) : Builder.getInt32(magics.s);
+    Q = Builder.CreateAShr(Q, MagicsShift);
+  }
+
+  // Extract the sign bit and add it to the quotient
+  Value *T = Builder.CreateLShr(Q, Shift);
+  Q = Builder.CreateAdd(Q, T);
+
+  return Q;
+}
+
+/// Generate code for unsigned division by a constant. Implementation follows
+/// TargetLowering::BuildUDIV that replaces division with multiplication
+/// by a magic number.
+static Value *generateUnsignedDivisionByConstant(Value *Dividend,
+                                                 ConstantInt *Divisor,
+                                                 IRBuilder<> &Builder) {
+  unsigned BitWidth = Dividend->getType()->getIntegerBitWidth();
+  assert((BitWidth == 64 || BitWidth == 32) && "Unexpected bit width");
+
+  APInt d = Divisor->getValue();
+  APInt::mu magics = d.magicu();
+
+  Value *Q = Dividend;
+
+  // If the divisor is even, we can avoid using the expensive fixup by shifting
+  // the divided value upfront.
+  if (magics.a != 0 && !d[0]) {
+    unsigned s = d.countTrailingZeros();
+    ConstantInt *Shift = (BitWidth == 64) ?
+      Builder.getInt64(s) : Builder.getInt32(s);
+    Q = Builder.CreateLShr(Q, Shift);
+
+    // Get magic number for the shifted divisor.
+    magics = d.lshr(s).magicu(s);
+    assert(magics.a == 0 && "Should use cheap fixup now");
+  }
+
+  // Multiply the numerator (operand 0) by the magic value
+  ConstantInt *Magics = (BitWidth == 64) ?
+    Builder.getInt64(magics.m.getZExtValue()) :
+    Builder.getInt32(magics.m.getZExtValue());
+  Q = Builder.CreateMul(Q, Magics);
+
+  if (magics.a == 0) {
+    assert(magics.s < d.getBitWidth() &&
+           "We shouldn't generate an undefined shift!");
+    ConstantInt *Shift = (BitWidth == 64) ?
+      Builder.getInt64(magics.s) : Builder.getInt32(magics.s);
+    Q = Builder.CreateLShr(Q, Shift);
+  } else {
+    Value *NPQ = Builder.CreateSub(Dividend, Q);
+    ConstantInt *One = (BitWidth == 64) ?
+      Builder.getInt64(1) : Builder.getInt32(1);
+    NPQ = Builder.CreateLShr(NPQ, One);
+    NPQ = Builder.CreateAdd(NPQ, Q);
+    ConstantInt *Shift = (BitWidth == 64) ?
+      Builder.getInt64(magics.s - 1) : Builder.getInt32(magics.s - 1);
+    Q = Builder.CreateLShr(NPQ, Shift);
+  }
+
+  return Q;
+}
+
 /// Generate code to calculate the remainder of two integers, replacing Rem with
 /// the generated code. This currently generates code using the udiv expansion,
 /// but future work includes generating more specialized code, e.g. when more
 /// information about the operands are known. Implements both 32bit and 64bit
 /// scalar division.
 ///
 /// @brief Replace Rem with generated code.
 bool llvm::expandRemainder(BinaryOperator *Rem) {
bool llvm::expandRemainder(BinaryOperator *Rem) {
   }
 
   return true;
 }
 
 
 /// Generate code to divide two integers, replacing Div with the generated
 /// code. This currently generates code similarly to compiler-rt's
-/// implementations, but future work includes generating more specialized code
-/// when more information about the operands are known. Implements both
-/// 32bit and 64bit scalar division.
+/// implementations. Specialized code is generated for division by a constant.
+/// Future work includes generating more specialized code when more information
+//  about the operands is known. Implements both 32bit and 64bit scalar division.
 ///
 /// @brief Replace Div with generated code.
 bool llvm::expandDivision(BinaryOperator *Div) {
   assert((Div->getOpcode() == Instruction::SDiv ||
           Div->getOpcode() == Instruction::UDiv) &&
          "Trying to expand division from a non-division function");
 
   IRBuilder<> Builder(Div);
 
   assert(!Div->getType()->isVectorTy() && "Div over vectors not supported");
   assert((Div->getType()->getIntegerBitWidth() == 32 ||
           Div->getType()->getIntegerBitWidth() == 64) &&
          "Div of bitwidth other than 32 or 64 not supported");
 
+  // handle a special case of dividing by a constant
+  if (ConstantInt *Divisor = dyn_cast<ConstantInt>(Div->getOperand(1))) {
+    Value *Quotient = NULL;
+    if (Div->getOpcode() == Instruction::SDiv) {
+      Quotient = generateSignedDivisionByConstant(Div->getOperand(0),
+                                                  Divisor, Builder);
+    } else {
+      Quotient = generateUnsignedDivisionByConstant(Div->getOperand(0),
+                                                    Divisor, Builder);
+    }
+
+    assert(Quotient && "failed to generate code for division by a constant");
+    Div->replaceAllUsesWith(Quotient);
+    Div->dropAllReferences();
+    Div->eraseFromParent();
+
+    return true;
+  }
+
   // First prepare the sign if it's a signed division
   if (Div->getOpcode() == Instruction::SDiv) {
     // Lower the code to unsigned division, and reset Div to point to the udiv.
     Value *Quotient = generateSignedDivisionCode(Div->getOperand(0),
                                                  Div->getOperand(1), Builder);
 
     // Check whether this is the insert point while Div is still valid.
     bool IsInsertPoint = Div->getIterator() == Builder.GetInsertPoint();

unittests/Transforms/Utils/IntegerDivision.cpp

TEST(IntegerDivision, URem64) {
 
   expandRemainder(cast<BinaryOperator>(Rem));
   EXPECT_TRUE(BB->front().getOpcode() == Instruction::ICmp);
 
   Instruction* Remainder = dyn_cast<Instruction>(cast<User>(Ret)->getOperand(0));
   EXPECT_TRUE(Remainder && Remainder->getOpcode() == Instruction::Sub);
 }
 
+TEST(IntegerDivision, SDivConst) {
+  LLVMContext C;
+  Module M("test division", C);
+  IRBuilder<> Builder(C);
+
+  Type *ArgTy1[] = { Builder.getInt32Ty() };
+  Function *F = Function::Create(FunctionType::get(Builder.getInt32Ty(),
+                                                   ArgTy1, false),
+                                 GlobalValue::ExternalLinkage, "F", &M);
+  assert(F->getArgumentList().size() == 1);
+
+  BasicBlock *BB = BasicBlock::Create(C, "", F);
+  Builder.SetInsertPoint(BB);
+
+  Function::arg_iterator AI = F->arg_begin();
+  Value *A = &*AI++;
+  Value *B = Builder.getInt32(13);
+
+  Value *Div = Builder.CreateSDiv(A, B);
+  EXPECT_TRUE(BB->front().getOpcode() == Instruction::SDiv);
+
+  Value *Ret = Builder.CreateRet(Div);
+
+  expandDivision(cast<BinaryOperator>(Div));
+  EXPECT_TRUE(BB->front().getOpcode() == Instruction::Mul);
+
+  Instruction* Quotient = dyn_cast<Instruction>(cast<User>(Ret)->getOperand(0));
+  EXPECT_TRUE(Quotient && Quotient->getOpcode() == Instruction::Add);
+}
+
+TEST(IntegerDivision, UDiv64Const) {
+  LLVMContext C;
+  Module M("test division", C);
+  IRBuilder<> Builder(C);
+
+  Type *ArgTy1[] = { Builder.getInt64Ty() };
+  Function *F = Function::Create(FunctionType::get(Builder.getInt64Ty(),
+                                                   ArgTy1, false),
+                                 GlobalValue::ExternalLinkage, "F", &M);
+  assert(F->getArgumentList().size() == 1);
+
+  BasicBlock *BB = BasicBlock::Create(C, "", F);
+  Builder.SetInsertPoint(BB);
+
+  Function::arg_iterator AI = F->arg_begin();
+  Value *A = &*AI++;
+  APInt d(64, 8);
+  ConstantInt *B = ConstantInt::get(C, d);
+  //Value *B = Builder.getInt64(8);
+  APInt::mu magics = d.magicu();
+
+  Value *Div = Builder.CreateUDiv(A, B);
+  EXPECT_TRUE(BB->front().getOpcode() == Instruction::UDiv);
+
+  Value *Ret = Builder.CreateRet(Div);
+
+  expandDivision(cast<BinaryOperator>(Div));
+
+  if (magics.a != 0 && !d[0]) {
+    EXPECT_TRUE(BB->front().getOpcode() == Instruction::LShr);
+  } else {
+    EXPECT_TRUE(BB->front().getOpcode() == Instruction::Mul);
+  }
+
+  Instruction* Quotient = dyn_cast<Instruction>(cast<User>(Ret)->getOperand(0));
+  EXPECT_TRUE(Quotient && Quotient->getOpcode() == Instruction::LShr);
+}
+
 }