Merge pull request google#226 from asraa:comb-to-cggi

PiperOrigin-RevId: 580164491

include/Dialect/Comb/IR/BUILD

@@ -15,7 +15,7 @@ td_library(
         "Comb.td",
         "Combinational.td",
     ],
-    includes = ["include"],
+    includes = ["../../../.."],
     deps = [
         "@llvm-project//mlir:BuiltinDialectTdFiles",
         "@llvm-project//mlir:ControlFlowInterfacesTdFiles",
@@ -45,7 +45,10 @@ gentbl_cc_library(
             "CombDialect.cpp.inc",
         ),
         (
-            ["-gen-dialect-doc"],
+            [
+                "-gen-dialect-doc",
+                "-dialect=comb",
+            ],
             "CombDialect.md",
         ),
     ],

include/Dialect/Comb/IR/Comb.td

@@ -37,6 +37,6 @@ def CombDialect : Dialect {
 class CombOp<string mnemonic, list<Trait> traits = []> :
     Op<CombDialect, mnemonic, traits>;
 
-include "Combinational.td"
+include "include/Dialect/Comb/IR/Combinational.td"
 
 #endif // HEIR_INCLUDE_DIALECT_COMB_COMB_TD

include/Dialect/Comb/IR/Combinational.td

@@ -14,7 +14,11 @@
 // Arithmetic and Logical Operations
 //===----------------------------------------------------------------------===//
 
+#ifndef HEIR_INCLUDE_DIALECT_COMB_COMBINATIONAL_TD
+#define HEIR_INCLUDE_DIALECT_COMB_COMBINATIONAL_TD
+
 include "mlir/Interfaces/InferTypeOpInterface.td"
+include "mlir/IR/BuiltinAttributes.td"
 include "mlir/IR/EnumAttr.td"
 
 def HWIntegerType : Type<
@@ -272,21 +276,22 @@ def TruthTableOp : CombOp<"truth_table", [Pure]> {
     ```
       %a = ... : i1
       %b = ... : i1
-      %0 = comb.truth_table %a, %b -> [false, true, true, false]
+      %0 = comb.truth_table %a, %b -> 6 : ui4
     ```
 
-    This operation assumes a fully elaborated table -- 2^n entries. Inputs are
-    sorted MSB -> LSB from left to right and the offset into `lookupTable` is
-    computed from them. The table is sorted from 0 -> (2^n - 1) from left to
-    right.
+    This operation assumes that the lookup table is described as an integer of
+    2^n bits to fully specify the table. Inputs are sorted MSB -> LSB from left
+    to right and the offset into `lookupTable` is computed from them. The
+    integer containing the truth table value's LSB is the output for the input
+    "all false", and the MSB is the output for the input "all true".
 
     No difference from array_get into an array of constants except for xprop
     behavior. If one of the inputs is unknown, but said input doesn't make a
     difference in the output (based on the lookup table) the result should not
     be 'x' -- it should be the well-known result.
   }];
 
-  let arguments = (ins Variadic<I1>:$inputs, BoolArrayAttr:$lookupTable);
+  let arguments = (ins Variadic<I1>:$inputs, Builtin_IntegerAttr:$lookupTable);
   let results = (outs I1:$result);
 
   let assemblyFormat = [{
@@ -295,3 +300,5 @@ def TruthTableOp : CombOp<"truth_table", [Pure]> {
 
   let hasVerifier = 1;
 }
+
+#endif // HEIR_INCLUDE_DIALECT_COMB_COMBINATIONAL_TD

lib/Dialect/Comb/IR/CombOps.cpp

@@ -221,9 +221,9 @@ LogicalResult TruthTableOp::verify() {
     return emitOpError("Truth tables support a maximum of ")
            << sizeof(size_t) * 8 - 1 << " inputs on your platform";
 
-  ArrayAttr table = getLookupTable();
-  if (table.size() != (1ull << numInputs))
-    return emitOpError("Expected lookup table of 2^n length");
+  auto table = getLookupTable();
+  if (table.getValue().getBitWidth() != (1ull << numInputs))
+    return emitOpError("Expected lookup table int of 2^n bits");
   return success();
 }
 

lib/Transforms/YosysOptimizer/LUTImporter.cpp

@@ -20,17 +20,19 @@ mlir::Operation *LUTImporter::createOp(Yosys::RTLIL::Cell *cell,
   assert(cell->type.begins_with("\\lut"));
 
   // Create truth table from cell attributes.
-  int lutSize;
-  StringRef(cell->type.substr(4, 1)).getAsInteger(10, lutSize);
-  SmallVector<bool> lutValues(1 << lutSize, false);
+  int lutBits;
+  StringRef(cell->type.substr(4, 1)).getAsInteger(10, lutBits);
 
-  for (int i = 0; i < lutValues.size(); i++) {
+  uint64_t lutValue = 0;
+  int lutSize = 1 << lutBits;
+  for (int i = 0; i < lutSize; i++) {
     auto lutStr =
         cell->getPort(Yosys::RTLIL::IdString(llvm::formatv("\\P{0}", i)));
-    lutValues[i] = lutStr.as_bool();
+    lutValue |= (lutStr.as_bool() ? 1 : 0) << i;
   }
 
-  auto lookupTable = b.getBoolArrayAttr(llvm::ArrayRef<bool>(lutValues));
+  auto lookupTable =
+      b.getIntegerAttr(b.getIntegerType(lutSize, /*isSigned=*/false), lutValue);
   return b.create<comb::TruthTableOp>(inputs, lookupTable);
 }
 

lib/Transforms/YosysOptimizer/LUTImporterTest.cpp

@@ -32,7 +32,6 @@ namespace mlir::heir {
 namespace {
 
 using bazel::tools::cpp::runfiles::Runfiles;
-using ::testing::ElementsAreArray;
 using ::testing::Test;
 
 static constexpr std::string_view kWorkspaceDir = "heir";
@@ -80,13 +79,8 @@ class LUTImporterTestFixture : public Test {
 // Note that we cannot lower truth tables to LLVM, so we must assert the IR
 // rather than executing the code.
 TEST_F(LUTImporterTestFixture, AddOneLUT3) {
-  std::vector<std::vector<bool>> expectedLuts = {
-      {0, 1, 1, 0, 0, 0, 0, 0}, {0, 0, 0, 1, 1, 1, 1, 0},
-      {0, 0, 0, 0, 0, 0, 0, 1}, {0, 1, 1, 0, 0, 0, 0, 0},
-      {0, 0, 0, 1, 1, 1, 1, 0}, {0, 0, 0, 0, 0, 0, 0, 1},
-      {0, 1, 1, 0, 0, 0, 0, 0}, {0, 0, 0, 1, 1, 1, 1, 0},
-      {0, 0, 0, 0, 0, 0, 0, 1}, {0, 1, 1, 0, 0, 0, 0, 0},
-      {1, 0, 0, 0, 0, 0, 0, 0}};
+  std::vector<uint8_t> expectedLuts = {6, 120, 128, 6, 120, 128,
+                                       6, 120, 128, 6, 1};
 
   auto func =
       runImporter("lib/Transforms/YosysOptimizer/tests/add_one_lut3.rtlil");
@@ -99,10 +93,8 @@ TEST_F(LUTImporterTestFixture, AddOneLUT3) {
 
   auto combOps = func.getOps<comb::TruthTableOp>().begin();
   for (size_t i = 0; i < expectedLuts.size(); i++) {
-    SmallVector<bool> table(llvm::map_range(
-        (*combOps++).getLookupTableAttr().getAsValueRange<IntegerAttr>(),
-        [](const APInt &a) { return !a.isZero(); }));
-    EXPECT_THAT(table, ElementsAreArray(expectedLuts[i]));
+    auto lutValue = (*combOps++).getLookupTable().getValue();
+    EXPECT_THAT(lutValue, APInt(lutValue.getBitWidth(), expectedLuts[i]));
   }
 }
 
@@ -118,10 +110,8 @@ TEST_F(LUTImporterTestFixture, AddOneLUT5) {
 
   auto combOps = func.getOps<comb::TruthTableOp>();
   for (auto combOp : combOps) {
-    SmallVector<bool> table(llvm::map_range(
-        combOp.getLookupTableAttr().getAsValueRange<IntegerAttr>(),
-        [](const APInt &a) { return !a.isZero(); }));
-    EXPECT_EQ(table.size(), 32);
+    auto lutValue = combOp.getLookupTable().getValue();
+    EXPECT_EQ(lutValue.getBitWidth(), 32);
   }
 }
 

tests/comb.mlir

@@ -2,7 +2,7 @@
 
 module {
     func.func @comb(%a: i1, %b: i1) -> () {
-        %0 = comb.truth_table %a, %b -> [true, false, true, false]
+        %0 = comb.truth_table %a, %b -> 6 : ui4
         return
     }
 }