[Moore] Power operator folders and canonicalizers (#7494)

llvm · Aug 9, 2024 · 09fc725 · 09fc725
1 parent d4e7857
commit 09fc725
Show file tree

Hide file tree

Showing 3 changed files with 101 additions and 0 deletions.
diff --git a/include/circt/Dialect/Moore/MooreOps.td b/include/circt/Dialect/Moore/MooreOps.td
@@ -686,6 +686,9 @@ class PowOpBase<string mnemonic> : BinaryOpBase<mnemonic> {
 
     See IEEE 1800-2017 § 11.4.3 "Arithmetic operators".
   }];
+
+  let hasCanonicalizeMethod = 1;
+  let hasFolder = 1;
 }
 
 def PowUOp : PowOpBase<"powu">;

diff --git a/lib/Dialect/Moore/MooreOps.cpp b/lib/Dialect/Moore/MooreOps.cpp
@@ -1101,6 +1101,74 @@ ReadOp::removeBlockingUses(const MemorySlot &slot,
   return DeletionKind::Delete;
 }
 
+//===----------------------------------------------------------------------===//
+// PowSOp
+//===----------------------------------------------------------------------===//
+
+static OpFoldResult powCommonFolding(MLIRContext *ctxt, Attribute lhs,
+                                     Attribute rhs) {
+  auto lhsValue = dyn_cast_or_null<FVIntegerAttr>(lhs);
+  if (lhsValue && lhsValue.getValue() == 1)
+    return lhs;
+
+  auto rhsValue = dyn_cast_or_null<FVIntegerAttr>(rhs);
+  if (rhsValue && rhsValue.getValue().isZero())
+    return FVIntegerAttr::get(ctxt,
+                              FVInt(rhsValue.getValue().getBitWidth(), 1));
+
+  return {};
+}
+
+OpFoldResult PowSOp::fold(FoldAdaptor adaptor) {
+  return powCommonFolding(getContext(), adaptor.getLhs(), adaptor.getRhs());
+}
+
+LogicalResult PowSOp::canonicalize(PowSOp op, PatternRewriter &rewriter) {
+  Location loc = op.getLoc();
+  auto intType = cast<IntType>(op.getRhs().getType());
+  if (auto baseOp = op.getLhs().getDefiningOp<ConstantOp>()) {
+    if (baseOp.getValue() == 2) {
+      Value constOne = rewriter.create<ConstantOp>(loc, intType, 1);
+      Value constZero = rewriter.create<ConstantOp>(loc, intType, 0);
+      Value shift = rewriter.create<ShlOp>(loc, constOne, op.getRhs());
+      Value isNegative = rewriter.create<SltOp>(loc, op.getRhs(), constZero);
+      auto condOp = rewriter.replaceOpWithNewOp<ConditionalOp>(
+          op, op.getLhs().getType(), isNegative);
+      Block *thenBlock = rewriter.createBlock(&condOp.getTrueRegion());
+      rewriter.setInsertionPointToStart(thenBlock);
+      rewriter.create<YieldOp>(loc, constZero);
+      Block *elseBlock = rewriter.createBlock(&condOp.getFalseRegion());
+      rewriter.setInsertionPointToStart(elseBlock);
+      rewriter.create<YieldOp>(loc, shift);
+      return success();
+    }
+  }
+
+  return failure();
+}
+
+//===----------------------------------------------------------------------===//
+// PowUOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult PowUOp::fold(FoldAdaptor adaptor) {
+  return powCommonFolding(getContext(), adaptor.getLhs(), adaptor.getRhs());
+}
+
+LogicalResult PowUOp::canonicalize(PowUOp op, PatternRewriter &rewriter) {
+  Location loc = op.getLoc();
+  auto intType = cast<IntType>(op.getRhs().getType());
+  if (auto baseOp = op.getLhs().getDefiningOp<ConstantOp>()) {
+    if (baseOp.getValue() == 2) {
+      Value constOne = rewriter.create<ConstantOp>(loc, intType, 1);
+      rewriter.replaceOpWithNewOp<ShlOp>(op, constOne, op.getRhs());
+      return success();
+    }
+  }
+
+  return failure();
+}
+
 //===----------------------------------------------------------------------===//
 // TableGen generated logic.
 //===----------------------------------------------------------------------===//

diff --git a/test/Dialect/Moore/canonicalizers.mlir b/test/Dialect/Moore/canonicalizers.mlir
@@ -250,3 +250,33 @@ func.func @ConvertConstantFourToTwoValued() -> (!moore.i42) {
   %1 = moore.conversion %0 : !moore.l42 -> !moore.i42
   return %1 : !moore.i42
 }
+
+// CHECK-LABEL: func @Pow
+func.func @Pow(%arg0 : !moore.l32) -> (!moore.l32, !moore.l32, !moore.l32, !moore.l32, !moore.l32, !moore.l32) {
+  // CHECK-NEXT: [[V0:%.+]] = moore.constant 1 : l32
+  // CHECK-NEXT: [[V1:%.+]] = moore.constant 0 : l32
+  %0 = moore.constant 0 : l32
+  %1 = moore.constant 1 : l32
+  %2 = moore.constant 2 : l32
+
+  %3 = moore.pows %1, %arg0 : l32
+  %4 = moore.pows %arg0, %0 : l32
+
+  %5 = moore.powu %1, %arg0 : l32
+  %6 = moore.powu %arg0, %0 : l32
+
+  // CHECK-NEXT: [[V2:%.+]] = moore.shl [[V0]], %arg0
+  // CHECK-NEXT: [[V3:%.+]] = moore.slt %arg0, [[V1]]
+  // CHECK-NEXT: [[V4:%.+]] = moore.conditional [[V3]]
+  // CHECK-NEXT:   moore.yield [[V1]]
+  // CHECK-NEXT: } {
+  // CHECK-NEXT:   moore.yield [[V2]]
+  // CHECK-NEXT: }
+  %7 = moore.pows %2, %arg0 : l32
+
+  // CHECK-NEXT: [[V5:%.+]] = moore.shl [[V0]], %arg0
+  %8 = moore.powu %2, %arg0 : l32
+
+  // CHECK-NEXT: return [[V0]], [[V0]], [[V0]], [[V0]], [[V4]], [[V5]] :
+  return %3, %4, %5, %6, %7, %8 : !moore.l32, !moore.l32, !moore.l32, !moore.l32, !moore.l32, !moore.l32
+}