[Comb] Canonicalize power-of-two unsigned div/mod. (#9177)

This commit adds canonicalization patterns for unsigned division and
modulo operations when the divisor is a power of two constant. These
operations can be efficiently lowered to bit extraction and
concatenation operations.

This is done in CombToSynth already but would be good to perform as part of canonicalization.

Author: uenoku

include/circt/Dialect/Comb/CombOps.h

@@ -90,6 +90,15 @@ Value createInject(OpBuilder &builder, Location loc, Value value,
 LogicalResult convertSubToAdd(comb::SubOp subOp,
                               mlir::PatternRewriter &rewriter);
 
+/// Convert unsigned division or modulo by a power of two.
+/// For division: divu(x, 2^n) -> concat(0...0, extract(x, n, width-n)).
+/// For modulo: modu(x, 2^n) -> concat(0...0, extract(x, 0, n))
+/// TODO: Support signed division and modulo.
+LogicalResult convertDivUByPowerOfTwo(DivUOp divOp,
+                                      mlir::PatternRewriter &rewriter);
+LogicalResult convertModUByPowerOfTwo(ModUOp modOp,
+                                      mlir::PatternRewriter &rewriter);
+
 /// Enum for mux chain folding styles.
 enum MuxChainWithComparisonFoldingStyle { None, BalancedMuxTree, ArrayGet };
 /// Mux chain folding that converts chains of muxes with index

include/circt/Dialect/Comb/Combinational.td

@@ -68,11 +68,11 @@ class UTVariadicOp<string mnemonic, list<Trait> traits = []> :
 def AddOp : UTVariadicOp<"add", [Commutative]>;
 def MulOp : UTVariadicOp<"mul", [Commutative]>;
 let hasFolder = true in {
-  def DivUOp : UTBinOp<"divu">;
   def DivSOp : UTBinOp<"divs">;
-  def ModUOp : UTBinOp<"modu">;
   def ModSOp : UTBinOp<"mods">;
   let hasCanonicalizeMethod = true in {
+    def DivUOp : UTBinOp<"divu">;
+    def ModUOp : UTBinOp<"modu">;
     def ShlOp : UTBinOp<"shl">;
     def ShrUOp : UTBinOp<"shru">;
     def ShrSOp : UTBinOp<"shrs">;

lib/Conversion/CombToSynth/CombToSynth.cpp

@@ -990,20 +990,8 @@ struct CombDivUOpConversion : DivModOpConversionBase<DivUOp> {
   matchAndRewrite(DivUOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     // Check if the divisor is a power of two.
-    if (auto rhsConstantOp = adaptor.getRhs().getDefiningOp<hw::ConstantOp>())
-      if (rhsConstantOp.getValue().isPowerOf2()) {
-        // Extract upper bits.
-        size_t extractAmount = rhsConstantOp.getValue().ceilLogBase2();
-        size_t width = op.getType().getIntOrFloatBitWidth();
-        Value upperBits = rewriter.createOrFold<comb::ExtractOp>(
-            op.getLoc(), adaptor.getLhs(), extractAmount,
-            width - extractAmount);
-        Value constZero = hw::ConstantOp::create(rewriter, op.getLoc(),
-                                                 APInt::getZero(extractAmount));
-        replaceOpWithNewOpAndCopyNamehint<comb::ConcatOp>(
-            rewriter, op, op.getType(), ArrayRef<Value>{constZero, upperBits});
-        return success();
-      }
+    if (llvm::succeeded(comb::convertDivUByPowerOfTwo(op, rewriter)))
+      return success();
 
     // When rhs is not power of two and the number of unknown bits are small,
     // create a mux tree that emulates all possible cases.
@@ -1024,19 +1012,8 @@ struct CombModUOpConversion : DivModOpConversionBase<ModUOp> {
   matchAndRewrite(ModUOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     // Check if the divisor is a power of two.
-    if (auto rhsConstantOp = adaptor.getRhs().getDefiningOp<hw::ConstantOp>())
-      if (rhsConstantOp.getValue().isPowerOf2()) {
-        // Extract lower bits.
-        size_t extractAmount = rhsConstantOp.getValue().ceilLogBase2();
-        size_t width = op.getType().getIntOrFloatBitWidth();
-        Value lowerBits = rewriter.createOrFold<comb::ExtractOp>(
-            op.getLoc(), adaptor.getLhs(), 0, extractAmount);
-        Value constZero = hw::ConstantOp::create(
-            rewriter, op.getLoc(), APInt::getZero(width - extractAmount));
-        replaceOpWithNewOpAndCopyNamehint<comb::ConcatOp>(
-            rewriter, op, op.getType(), ArrayRef<Value>{constZero, lowerBits});
-        return success();
-      }
+    if (llvm::succeeded(comb::convertModUByPowerOfTwo(op, rewriter)))
+      return success();
 
     // When rhs is not power of two and the number of unknown bits are small,
     // create a mux tree that emulates all possible cases.

lib/Dialect/Comb/CombFolds.cpp

@@ -1697,6 +1697,20 @@ OpFoldResult ModSOp::fold(FoldAdaptor adaptor) {
     return {};
   return foldMod<ModSOp, /*isSigned=*/true>(*this, adaptor.getOperands());
 }
+
+LogicalResult DivUOp::canonicalize(DivUOp op, PatternRewriter &rewriter) {
+  if (isOpTriviallyRecursive(op) || !op.getTwoState())
+    return failure();
+  return convertDivUByPowerOfTwo(op, rewriter);
+}
+
+LogicalResult ModUOp::canonicalize(ModUOp op, PatternRewriter &rewriter) {
+  if (isOpTriviallyRecursive(op) || !op.getTwoState())
+    return failure();
+
+  return convertModUByPowerOfTwo(op, rewriter);
+}
+
 //===----------------------------------------------------------------------===//
 // ConcatOp
 //===----------------------------------------------------------------------===//

lib/Dialect/Comb/CombOps.cpp

@@ -234,6 +234,72 @@ llvm::LogicalResult comb::convertSubToAdd(comb::SubOp subOp,
   return success();
 }
 
+static llvm::LogicalResult convertDivModUByPowerOfTwo(PatternRewriter &rewriter,
+                                                      Operation *op, Value lhs,
+                                                      Value rhs, bool isDiv) {
+  // Check if the divisor is a power of two constant.
+  auto rhsConstantOp = rhs.getDefiningOp<hw::ConstantOp>();
+  if (!rhsConstantOp)
+    return failure();
+
+  APInt rhsValue = rhsConstantOp.getValue();
+  if (!rhsValue.isPowerOf2())
+    return failure();
+
+  Location loc = op->getLoc();
+
+  unsigned width = lhs.getType().getIntOrFloatBitWidth();
+  unsigned bitPosition = rhsValue.ceilLogBase2();
+
+  if (isDiv) {
+    // divu(x, 2^n) -> concat(0...0, extract(x, n, width-n))
+    // This is equivalent to a right shift by n bits.
+
+    // Extract the upper bits (equivalent to right shift).
+    Value upperBits = rewriter.createOrFold<comb::ExtractOp>(
+        loc, lhs, bitPosition, width - bitPosition);
+
+    // Concatenate with zeros on the left.
+    Value zeros =
+        hw::ConstantOp::create(rewriter, loc, APInt::getZero(bitPosition));
+
+    // use replaceOpWithNewOpAndCopyNamehint?
+    replaceOpAndCopyNamehint(
+        rewriter, op,
+        comb::ConcatOp::create(rewriter, loc,
+                               ArrayRef<Value>{zeros, upperBits}));
+    return success();
+  }
+
+  // modu(x, 2^n) -> concat(0...0, extract(x, 0, n))
+  // This extracts the lower n bits (equivalent to bitwise AND with 2^n - 1).
+
+  // Extract the lower bits.
+  Value lowerBits =
+      rewriter.createOrFold<comb::ExtractOp>(loc, lhs, 0, bitPosition);
+
+  // Concatenate with zeros on the left.
+  Value zeros = hw::ConstantOp::create(rewriter, loc,
+                                       APInt::getZero(width - bitPosition));
+
+  replaceOpAndCopyNamehint(
+      rewriter, op,
+      comb::ConcatOp::create(rewriter, loc, ArrayRef<Value>{zeros, lowerBits}));
+  return success();
+}
+
+LogicalResult comb::convertDivUByPowerOfTwo(DivUOp divOp,
+                                            mlir::PatternRewriter &rewriter) {
+  return convertDivModUByPowerOfTwo(rewriter, divOp, divOp.getLhs(),
+                                    divOp.getRhs(), /*isDiv=*/true);
+}
+
+LogicalResult comb::convertModUByPowerOfTwo(ModUOp modOp,
+                                            mlir::PatternRewriter &rewriter) {
+  return convertDivModUByPowerOfTwo(rewriter, modOp, modOp.getLhs(),
+                                    modOp.getRhs(), /*isDiv=*/false);
+}
+
 //===----------------------------------------------------------------------===//
 // ICmpOp
 //===----------------------------------------------------------------------===//

lib/Dialect/Datapath/DatapathFolds.cpp

@@ -7,10 +7,12 @@
 //===----------------------------------------------------------------------===//
 
 #include "circt/Dialect/Comb/CombOps.h"
+#include "circt/Dialect/Datapath/DatapathDialect.h"
 #include "circt/Dialect/Datapath/DatapathOps.h"
 #include "circt/Dialect/HW/HWOps.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
+#include "llvm/Support/Casting.h"
 #include "llvm/Support/KnownBits.h"
 #include <algorithm>
 
@@ -117,6 +119,9 @@ struct FoldAddIntoCompress : public OpRewritePattern<comb::AddOp> {
   using OpRewritePattern::OpRewritePattern;
 
   // add(compress(a,b,c),d) -> add(compress(a,b,c,d))
+  // FIXME: This should be implemented as a canonicalization pattern for
+  // compress op. Currently `hasDatapathOperand` flag prevents introducing
+  // datapath operations from comb operations.
   LogicalResult matchAndRewrite(comb::AddOp addOp,
                                 PatternRewriter &rewriter) const override {
     // comb.add canonicalization patterns handle folding add operations
@@ -128,15 +133,23 @@ struct FoldAddIntoCompress : public OpRewritePattern<comb::AddOp> {
     llvm::SmallSetVector<Value, 8> processedCompressorResults;
     SmallVector<Value, 8> newCompressOperands;
     // Only construct compressor if can form a larger compressor than what
-    // is currently an input of this add
-    bool shouldFold = false;
+    // is currently an input of this add. Also check that there is at least
+    // one datapath operand.
+    bool shouldFold = false, hasDatapathOperand = false;
 
     for (Value operand : operands) {
 
       // Skip if already processed this compressor
       if (processedCompressorResults.contains(operand))
         continue;
 
+      if (auto *op = operand.getDefiningOp()) {
+        if (llvm::isa_and_nonnull<datapath::DatapathDialect>(
+                op->getDialect())) {
+          hasDatapathOperand = true;
+        }
+      }
+
       // If the operand has multiple uses, we do not fold it into a compress
       // operation, so we treat it as a regular operand.
       if (!operand.hasOneUse()) {
@@ -174,7 +187,7 @@ struct FoldAddIntoCompress : public OpRewritePattern<comb::AddOp> {
 
     // Only fold if we have constructed a larger compressor than what was
     // already there
-    if (!shouldFold)
+    if (!shouldFold || !hasDatapathOperand)
       return failure();
 
     // Create a new CompressOp with all collected operands

lib/Dialect/Synth/Transforms/SynthesisPipeline.cpp

@@ -59,8 +59,9 @@ void circt::synth::buildCombLoweringPipeline(
       circt::ConvertDatapathToCombOptions datapathOptions;
       datapathOptions.timingAware = options.timingAware;
       pm.addPass(createConvertDatapathToComb(datapathOptions));
-      pm.addPass(createSimpleCanonicalizerPass());
     }
+    pm.addPass(createCSEPass());
+    pm.addPass(createSimpleCanonicalizerPass());
     // Partially legalize Comb, then run CSE and canonicalization.
     circt::ConvertCombToSynthOptions convOptions;
     addOpName<comb::AndOp, comb::OrOp, comb::XorOp, comb::MuxOp, comb::ICmpOp,

test/Dialect/Comb/canonicalization.mlir

@@ -1240,6 +1240,71 @@ hw.module @moduloZeroDividend(in %arg0 : i32, out o1: i32, out o2: i32) {
   hw.output %0, %1 : i32, i32
 }
 
+// CHECK-LABEL: hw.module @divuPowerOfTwo
+hw.module @divuPowerOfTwo(in %arg0 : i8, out o1: i8, out o2: i8, out o3: i8, out o4: i8, out o5: i8) {
+  // divu(x, 2) -> concat(0, extract(x, 1, 7))
+  // CHECK: [[EXT1:%.+]] = comb.extract %arg0 from 1 : (i8) -> i7
+  // CHECK-NEXT: [[RES1:%.+]] = comb.concat %false, [[EXT1]] : i1, i7
+  %c2 = hw.constant 2 : i8
+  %0 = comb.divu bin %arg0, %c2 : i8
+
+  // divu(x, 4) -> concat(00, extract(x, 2, 6))
+  // CHECK-NEXT: [[EXT2:%.+]] = comb.extract %arg0 from 2 : (i8) -> i6
+  // CHECK-NEXT: [[RES2:%.+]] = comb.concat %c0_i2, [[EXT2]] : i2, i6
+  %c4 = hw.constant 4 : i8
+  %1 = comb.divu bin %arg0, %c4 : i8
+
+  // divu(x, 16) -> concat(0000, extract(x, 4, 4))
+  // CHECK-NEXT: [[EXT4:%.+]] = comb.extract %arg0 from 4 : (i8) -> i4
+  // CHECK-NEXT: [[RES3:%.+]] = comb.concat %c0_i4, [[EXT4]] : i4, i4
+  %c16 = hw.constant 16 : i8
+  %2 = comb.divu bin %arg0, %c16 : i8
+
+  // divu(x, 3) -> not canonicalized (not power of two)
+  // CHECK-NEXT: [[RES4:%.+]] = comb.divu bin %arg0, %c3_i8 : i8
+  %c3 = hw.constant 3 : i8
+  %3 = comb.divu bin %arg0, %c3 : i8
+
+  // Make sure canonicalization does not happen if there is no bin flag.
+  // CHECK-NEXT: [[RES5:%.+]] = comb.divu %arg0, %c2_i8 : i8
+  %4 = comb.divu %arg0, %c2 : i8
+
+  // CHECK: hw.output [[RES1]], [[RES2]], [[RES3]], [[RES4]], [[RES5]]
+  hw.output %0, %1, %2, %3, %4 : i8, i8, i8, i8, i8
+}
+
+// CHECK-LABEL: hw.module @moduPowerOfTwo
+hw.module @moduPowerOfTwo(in %arg0 : i8, out o1: i8, out o2: i8, out o3: i8, out o4: i8, out o5: i8) {
+  // modu(x, 2) -> concat(0000000, extract(x, 0, 1))
+  // CHECK:      [[EXT1:%.+]] = comb.extract %arg0 from 0 : (i8) -> i1
+  // CHECK-NEXT: [[RES1:%.+]] = comb.concat %c0_i7, [[EXT1]] : i7, i1
+  %c2 = hw.constant 2 : i8
+  %0 = comb.modu bin %arg0, %c2 : i8
+
+  // modu(x, 4) -> concat(000000, extract(x, 0, 2))
+  // CHECK-NEXT: [[EXT2:%.+]] = comb.extract %arg0 from 0 : (i8) -> i2
+  // CHECK-NEXT: [[RES2:%.+]] = comb.concat %c0_i6, [[EXT2]] : i6, i2
+  %c4 = hw.constant 4 : i8
+  %1 = comb.modu bin %arg0, %c4 : i8
+
+  // modu(x, 16) -> concat(0000, extract(x, 0, 4))
+  // CHECK-NEXT: [[EXT4:%.+]] = comb.extract %arg0 from 0 : (i8) -> i4
+  // CHECK-NEXT: [[RES3:%.+]] = comb.concat %c0_i4, [[EXT4]] : i4, i4
+  %c16 = hw.constant 16 : i8
+  %2 = comb.modu bin %arg0, %c16 : i8
+
+  // modu(x, 3) -> not canonicalized (not power of two)
+  // CHECK-NEXT: [[RES4:%.+]] = comb.modu bin %arg0, %c3_i8 : i8
+  %c3 = hw.constant 3 : i8
+  %3 = comb.modu bin %arg0, %c3 : i8
+
+  // Make sure canonicalization does not happen if there is no bin flag.
+  // CHECK-NEXT: [[RES5:%.+]] = comb.modu %arg0, %c2_i8 : i8
+  %4 = comb.modu %arg0, %c2 : i8
+  // CHECK: hw.output [[RES1]], [[RES2]], [[RES3]], [[RES4]], [[RES5]]
+  hw.output %0, %1, %2, %3, %4 : i8, i8, i8, i8, i8
+}
+
 // CHECK-LABEL: hw.module @orWithNegation
 hw.module @orWithNegation(in %arg0 : i32, out o1: i32) {
   // CHECK: [[ALLONES:%.*]] = hw.constant -1 : i32

test/Dialect/Datapath/canonicalization.mlir

@@ -115,3 +115,13 @@ hw.module @pos_partial_product_do_nothing(in %a : i4, in %b : i4, in %c : i4, ou
   %2:4 = datapath.partial_product %0, %1 : (i4, i4) -> (i4, i4, i4, i4)
   hw.output %2#0, %2#1, %2#2, %2#3 : i4, i4, i4, i4
 }
+
+// CHECK-LABEL: @dont_introduce_compressor
+hw.module @dont_introduce_compressor(in %a : i4, in %b : i4, in %c: i4, out sum : i4) {
+  // CHECK-NOT: datapath.compress
+  // CHECK-NEXT: comb.add
+  // CHECK-NEXT: hw.output
+  %0:4 = datapath.partial_product %a, %b : (i4, i4) -> (i4, i4, i4, i4)
+  %1 = comb.add %a, %b, %c : i4
+  hw.output %1 : i4
+}