add a ModArith constant op

This custom-parsed op and attribute make it relatively easy to systematically change an `arith.constant` op to `mod_arith`: ``` arith.constant 17 : i32 --> mod_arith.constant 17 : !mod_arith.mod_arith<234 : i32> ``` - comment on getValue nightmare - mod_arith.mod_arith -> mod_arith.int
google · Nov 19, 2024 · 10721ef · 10721ef
1 parent 6d91164
commit 10721ef
Show file tree

Hide file tree

Showing 17 changed files with 198 additions and 28 deletions.
diff --git a/lib/Dialect/ModArith/Conversions/ModArithToArith/ModArithToArith.cpp b/lib/Dialect/ModArith/Conversions/ModArithToArith/ModArithToArith.cpp
@@ -39,8 +39,8 @@ class ModArithToArithTypeConverter : public TypeConverter {
   }
 };
 
-// A herlper function to generate the attribute or type
-// needed to represent the result of modarith op as an integer
+// A helper function to generate the attribute or type
+// needed to represent the result of mod_arith op as an integer
 // before applying a remainder operation
 template <typename Op>
 TypedAttr modulusAttr(Op op, bool mul = false) {

diff --git a/lib/Dialect/ModArith/IR/BUILD b/lib/Dialect/ModArith/IR/BUILD
@@ -13,11 +13,13 @@ cc_library(
         "ModArithDialect.cpp",
     ],
     hdrs = [
+        "ModArithAttributes.h",
         "ModArithDialect.h",
         "ModArithOps.h",
         "ModArithTypes.h",
     ],
     deps = [
+        ":attributes_inc_gen",
         ":dialect_inc_gen",
         ":ops_inc_gen",
         ":types_inc_gen",
@@ -32,6 +34,7 @@ cc_library(
 td_library(
     name = "td_files",
     srcs = [
+        "ModArithAttributes.td",
         "ModArithDialect.td",
         "ModArithOps.td",
         "ModArithTypes.td",
@@ -70,6 +73,36 @@ gentbl_cc_library(
     ],
 )
 
+gentbl_cc_library(
+    name = "attributes_inc_gen",
+    tbl_outs = [
+        (
+            [
+                "-gen-attrdef-decls",
+                "-attrdefs-dialect=mod_arith",
+            ],
+            "ModArithAttributes.h.inc",
+        ),
+        (
+            [
+                "-gen-attrdef-defs",
+                "-attrdefs-dialect=mod_arith",
+            ],
+            "ModArithAttributes.cpp.inc",
+        ),
+        (
+            ["-gen-attrdef-doc"],
+            "ModArithAttributes.md",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "ModArithAttributes.td",
+    deps = [
+        ":dialect_inc_gen",
+        ":td_files",
+    ],
+)
+
 gentbl_cc_library(
     name = "types_inc_gen",
     tbl_outs = [

diff --git a/lib/Dialect/ModArith/IR/ModArithAttributes.h b/lib/Dialect/ModArith/IR/ModArithAttributes.h
@@ -0,0 +1,10 @@
+#ifndef LIB_DIALECT_MODARITH_IR_MODARITHATTRIBUTES_H_
+#define LIB_DIALECT_MODARITH_IR_MODARITHATTRIBUTES_H_
+
+#include "lib/Dialect/ModArith/IR/ModArithDialect.h"
+#include "lib/Dialect/ModArith/IR/ModArithTypes.h"
+
+#define GET_ATTRDEF_CLASSES
+#include "lib/Dialect/ModArith/IR/ModArithAttributes.h.inc"
+
+#endif  // LIB_DIALECT_MODARITH_IR_MODARITHATTRIBUTES_H_
diff --git a/lib/Dialect/ModArith/IR/ModArithAttributes.td b/lib/Dialect/ModArith/IR/ModArithAttributes.td
@@ -0,0 +1,43 @@
+#ifndef LIB_DIALECT_MODARITH_IR_MODARITHATTRS_TD_
+#define LIB_DIALECT_MODARITH_IR_MODARITHATTRS_TD_
+
+include "lib/Dialect/ModArith/IR/ModArithDialect.td"
+include "lib/Dialect/ModArith/IR/ModArithTypes.td"
+include "mlir/IR/BuiltinAttributes.td"
+
+
+class ModArith_Attr<string name, string attrMnemonic, list<Trait> traits = []>
+  : AttrDef<ModArith_Dialect, name, traits> {
+  let mnemonic = attrMnemonic;
+}
+
+
+def ModArith_ModArithAttr : ModArith_Attr<
+    "ModArith", "int", [TypedAttrInterface]> {
+  let summary = "a typed mod_arith attribute";
+  let description = [{
+    Example:
+
+    ```mlir
+    #attr = 123:i32
+    #attr_verbose = #mod_arith.int<123:i32>
+    ```
+  }];
+  let parameters = (ins "::mlir::heir::mod_arith::ModArithType":$type, "mlir::IntegerAttr":$value);
+  let assemblyFormat = "$value `:` $type";
+  let builders = [
+    AttrBuilderWithInferredContext<(ins "ModArithType":$type,
+                                        "const int64_t":$value), [{
+      return $_get(
+        type.getContext(),
+        type,
+        IntegerAttr::get(type.getModulus().getType(), value));
+    }]>
+  ];
+  let extraClassDeclaration = [{
+    using ValueType = ::mlir::Attribute;
+  }];
+}
+
+
+#endif  // LIB_DIALECT_MODARITH_IR_MODARITHATTRS_TD_
diff --git a/lib/Dialect/ModArith/IR/ModArithDialect.cpp b/lib/Dialect/ModArith/IR/ModArithDialect.cpp
@@ -10,7 +10,8 @@
 #include "mlir/include/mlir/Support/LogicalResult.h"     // from @llvm-project
 
 // NOLINTBEGIN(misc-include-cleaner): Required to define ModArithDialect,
-// ModArithTypes, ModArithOps
+// ModArithTypes, ModArithOps, ModArithAttributes
+#include "lib/Dialect/ModArith/IR/ModArithAttributes.h"
 #include "lib/Dialect/ModArith/IR/ModArithOps.h"
 #include "lib/Dialect/ModArith/IR/ModArithTypes.h"
 #include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
@@ -19,6 +20,9 @@
 // Generated definitions
 #include "lib/Dialect/ModArith/IR/ModArithDialect.cpp.inc"
 
+#define GET_ATTRDEF_CLASSES
+#include "lib/Dialect/ModArith/IR/ModArithAttributes.cpp.inc"
+
 #define GET_TYPEDEF_CLASSES
 #include "lib/Dialect/ModArith/IR/ModArithTypes.cpp.inc"
 
@@ -34,6 +38,10 @@ void ModArithDialect::initialize() {
 #define GET_TYPEDEF_LIST
 #include "lib/Dialect/ModArith/IR/ModArithTypes.cpp.inc"
       >();
+  addAttributes<
+#define GET_ATTRDEF_LIST
+#include "lib/Dialect/ModArith/IR/ModArithAttributes.cpp.inc"
+      >();
   addOperations<
 #define GET_OP_LIST
 #include "lib/Dialect/ModArith/IR/ModArithOps.cpp.inc"
@@ -128,6 +136,54 @@ LogicalResult BarrettReduceOp::verify() {
   return success();
 }
 
+ParseResult ConstantOp::parse(OpAsmParser &parser, OperationState &result) {
+  APInt parsedValue(64, 0);
+  Type parsedType;
+
+  if (failed(parser.parseInteger(parsedValue))) {
+    parser.emitError(parser.getCurrentLocation(),
+                     "found invalid integer value");
+    return failure();
+  }
+
+  if (parser.parseColon() || parser.parseType(parsedType)) return failure();
+
+  auto modArithType = dyn_cast<ModArithType>(parsedType);
+  if (!modArithType) return failure();
+
+  auto outputBitWidth =
+      modArithType.getModulus().getType().getIntOrFloatBitWidth();
+  if (parsedValue.getActiveBits() > outputBitWidth)
+    return parser.emitError(parser.getCurrentLocation(),
+                            "constant value is too large for the modulus");
+
+  auto intValue = IntegerAttr::get(modArithType.getModulus().getType(),
+                                   parsedValue.trunc(outputBitWidth));
+  result.addAttribute(
+      "value", ModArithAttr::get(parser.getContext(), modArithType, intValue));
+  result.addTypes(modArithType);
+  return success();
+}
+
+void ConstantOp::print(OpAsmPrinter &p) {
+  p << " ";
+  // getValue chain:
+  // op's ModArithAttribute value
+  //   -> ModArithAttribute's IntegerAttr value
+  //   -> IntegerAttr's APInt value
+  getValue().getValue().getValue().print(p.getStream(), true);
+  p << " : ";
+  p.printType(getOutput().getType());
+}
+
+LogicalResult ConstantOp::inferReturnTypes(
+    mlir::MLIRContext *context, std::optional<mlir::Location> loc,
+    ConstantOpAdaptor adaptor, llvm::SmallVectorImpl<mlir::Type> &returnTypes) {
+  adaptor.getValue().dump();
+  returnTypes.push_back(adaptor.getValue().getType());
+  return success();
+}
+
 }  // namespace mod_arith
 }  // namespace heir
 }  // namespace mlir
diff --git a/lib/Dialect/ModArith/IR/ModArithDialect.td b/lib/Dialect/ModArith/IR/ModArithDialect.td
@@ -11,6 +11,7 @@ def ModArith_Dialect : Dialect {
 
   let cppNamespace = "::mlir::heir::mod_arith";
   let useDefaultTypePrinterParser = 1;
+  let useDefaultAttributePrinterParser = 1;
 
   let dependentDialects = [
     "arith::ArithDialect",

diff --git a/lib/Dialect/ModArith/IR/ModArithOps.h b/lib/Dialect/ModArith/IR/ModArithOps.h
@@ -2,6 +2,7 @@
 #define LIB_DIALECT_MODARITH_IR_MODARITHOPS_H_
 
 // NOLINTBEGIN(misc-include-cleaner): Required to define ModArithOps
+#include "lib/Dialect/ModArith/IR/ModArithAttributes.h"
 #include "lib/Dialect/ModArith/IR/ModArithDialect.h"
 #include "lib/Dialect/ModArith/IR/ModArithTypes.h"
 #include "mlir/include/mlir/IR/BuiltinOps.h"  // from @llvm-project

diff --git a/lib/Dialect/ModArith/IR/ModArithOps.td b/lib/Dialect/ModArith/IR/ModArithOps.td
@@ -2,6 +2,7 @@
 #define LIB_DIALECT_MODARITH_IR_MODARITHOPS_TD_
 
 include "lib/Dialect/ModArith/IR/ModArithDialect.td"
+include "lib/Dialect/ModArith/IR/ModArithAttributes.td"
 include "lib/Dialect/ModArith/IR/ModArithTypes.td"
 include "mlir/IR/BuiltinAttributes.td"
 include "mlir/IR/CommonTypeConstraints.td"
@@ -28,8 +29,8 @@ def ModArith_EncapsulateOp : ModArith_Op<"encapsulate", [Pure, ElementwiseMappab
 
     Examples:
     ```
-    mod_arith.encapsulate %c0 : i32 -> mod_arith.mod_arith<65537 : i32>
-    mod_arith.encapsulate %c1 : i64 -> mod_arith.mod_arith<65537>
+    mod_arith.encapsulate %c0 : i32 -> mod_arith.int<65537 : i32>
+    mod_arith.encapsulate %c1 : i64 -> mod_arith.int<65537>
     ```
   }];
 
@@ -52,10 +53,10 @@ def ModArith_ExtractOp : ModArith_Op<"extract", [Pure, ElementwiseMappable]> {
 
     Examples:
     ```
-    %m0 = mod_arith.encapsulate %c0 : i32 -> mod_arith.mod_arith<65537 : i32>
-    %m1 = mod_arith.encapsulate %c1 : i64 -> mod_arith.mod_arith<65537>
-    %c2 = mod_arith.extract %m0 : mod_arith.mod_arith<65537 : i32> -> i32
-    %c3 = mod_arith.extract %m1 : mod_arith.mod_arith<65537> -> i64
+    %m0 = mod_arith.encapsulate %c0 : i32 -> mod_arith.int<65537 : i32>
+    %m1 = mod_arith.encapsulate %c1 : i64 -> mod_arith.int<65537>
+    %c2 = mod_arith.extract %m0 : mod_arith.int<65537 : i32> -> i32
+    %c3 = mod_arith.extract %m1 : mod_arith.int<65537> -> i64
     ```
   }];
 
@@ -67,6 +68,28 @@ def ModArith_ExtractOp : ModArith_Op<"extract", [Pure, ElementwiseMappable]> {
   let assemblyFormat = "operands attr-dict `:` type($input) `->` type($output)";
 }
 
+def ModArith_ConstantOp : Op<ModArith_Dialect, "constant",
+    [Pure, InferTypeOpAdaptor]> {
+  let summary = "Define a constant value via an attribute.";
+  let description = [{
+    Example:
+
+    ```mlir
+    %0 = mod_arith.constant 123 : !mod_arith.int<65537:i32>
+    ```
+  }];
+  let arguments = (ins ModArith_ModArithAttr:$value);
+  let results = (outs ModArith_ModArithType:$output);
+  let hasCustomAssemblyFormat = 1;
+
+  let builders = [
+    OpBuilder<(ins "::mlir::heir::mod_arith::ModArithType":$ty, "int64_t":$value), [{
+      return build($_builder, $_state, ty, mod_arith::ModArithAttr::get(ty, value));
+    }]>
+  ];
+}
+
+
 def ModArith_ReduceOp : ModArith_Op<"reduce", [Pure, ElementwiseMappable, SameOperandsAndResultType]> {
   let summary = "reduce the mod arith type to its canonical representative";
 
@@ -77,9 +100,9 @@ def ModArith_ReduceOp : ModArith_Op<"reduce", [Pure, ElementwiseMappable, SameOp
     Examples:
     ```
     %c0 = arith.constant 65538 : i32
-    %m0 = mod_arith.encapsulate %c0 : i32 -> mod_arith.mod_arith<65537 : i32>
+    %m0 = mod_arith.encapsulate %c0 : i32 -> mod_arith.int<65537 : i32>
     // mod_arith.extract %m0 produces 65538
-    %m1 = mod_arith.reduce %m0 : mod_arith.mod_arith<65537: i32>
+    %m1 = mod_arith.reduce %m0 : mod_arith.int<65537: i32>
     // mod_arith.extract %m1 produces 1
     ```
   }];

diff --git a/lib/Dialect/ModArith/IR/ModArithTypes.td b/lib/Dialect/ModArith/IR/ModArithTypes.td
@@ -12,16 +12,16 @@ class ModArith_Type<string name, string typeMnemonic>
   let mnemonic = typeMnemonic;
 }
 
-def ModArith_ModArith : ModArith_Type<"ModArith", "mod_arith"> {
+def ModArith_ModArithType : ModArith_Type<"ModArith", "int"> {
   let summary = "Integer type with modular arithmetic";
   let description = [{
-    `mod_arith.mod_arith<p>` represents an element of the ring of integers modulo $p$.
+    `mod_arith.int<p>` represents an element of the ring of integers modulo $p$.
     The `modulus` attribute is the ring modulus, and `mod_arith` operations lower to
     `arith` operations that produce results in the range `[0, modulus)`, often called
     the _canonical representative_.
 
     `modulus` is specified with an integer type suffix, for example,
-    `mod_arith.mod_arith<65537 : i32>`. This corresponds to the storage type for the
+    `mod_arith.int<65537 : i32>`. This corresponds to the storage type for the
     modulus, and is `i64` by default.
 
     It is required that the underlying integer type should be larger than
@@ -40,9 +40,9 @@ def ModArith_ModArith : ModArith_Type<"ModArith", "mod_arith"> {
 
     Examples:
     ```
-    !Zp1 = !mod_arith.mod_arith<7> // implicitly being i64
-    !Zp2 = !mod_arith.mod_arith<65537 : i32>
-    !Zp3 = !mod_arith.mod_arith<536903681 : i64>
+    !Zp1 = !mod_arith.int<7> // implicitly being i64
+    !Zp2 = !mod_arith.int<65537 : i32>
+    !Zp3 = !mod_arith.int<536903681 : i64>
     ```
   }];
   let parameters = (ins
@@ -51,6 +51,6 @@ def ModArith_ModArith : ModArith_Type<"ModArith", "mod_arith"> {
   let assemblyFormat = "`<` $modulus `>`";
 }
 
-def ModArithLike: TypeOrContainer<ModArith_ModArith, "mod_arith-like">;
+def ModArithLike: TypeOrContainer<ModArith_ModArithType, "mod_arith-like">;
 
 #endif  // LIB_TYPES_MODARITH_IR_MODARITHTYPES_TD_
diff --git a/tests/Dialect/ModArith/Conversions/mod_arith_to_arith/mod-arith-to-arith.mlir b/tests/Dialect/ModArith/Conversions/mod_arith_to_arith/mod-arith-to-arith.mlir
@@ -1,6 +1,6 @@
 // RUN: heir-opt -mod-arith-to-arith --split-input-file %s | FileCheck %s --enable-var-scope
 
-!Zp = !mod_arith.mod_arith<65537 : i32>
+!Zp = !mod_arith.int<65537 : i32>
 !Zpv = tensor<4x!Zp>
 
 // CHECK-LABEL: @test_lower_encapsulate

diff --git a/tests/Dialect/ModArith/Conversions/mod_arith_to_arith/runner/lower_add.mlir b/tests/Dialect/ModArith/Conversions/mod_arith_to_arith/runner/lower_add.mlir
@@ -5,7 +5,7 @@
 
 func.func private @printMemrefI32(memref<*xi32>) attributes { llvm.emit_c_interface }
 
-!Zp = !mod_arith.mod_arith<7681 : i26>
+!Zp = !mod_arith.int<7681 : i26>
 !Zpv = tensor<4x!Zp>
 
 func.func @test_lower_add() {

diff --git a/tests/Dialect/ModArith/Conversions/mod_arith_to_arith/runner/lower_mac.mlir b/tests/Dialect/ModArith/Conversions/mod_arith_to_arith/runner/lower_mac.mlir
@@ -5,7 +5,7 @@
 
 func.func private @printMemrefI32(memref<*xi32>) attributes { llvm.emit_c_interface }
 
-!Zp = !mod_arith.mod_arith<7681 : i26>
+!Zp = !mod_arith.int<7681 : i26>
 !Zpv = tensor<4x!Zp>
 
 func.func @test_lower_mac() {

diff --git a/tests/Dialect/ModArith/Conversions/mod_arith_to_arith/runner/lower_mul.mlir b/tests/Dialect/ModArith/Conversions/mod_arith_to_arith/runner/lower_mul.mlir
@@ -5,7 +5,7 @@
 
 func.func private @printMemrefI32(memref<*xi32>) attributes { llvm.emit_c_interface }
 
-!Zp = !mod_arith.mod_arith<7681 : i26>
+!Zp = !mod_arith.int<7681 : i26>
 !Zpv = tensor<4x!Zp>
 
 func.func @test_lower_mul() {

diff --git a/tests/Dialect/ModArith/Conversions/mod_arith_to_arith/runner/lower_reduce.mlir b/tests/Dialect/ModArith/Conversions/mod_arith_to_arith/runner/lower_reduce.mlir
@@ -5,10 +5,10 @@
 
 func.func private @printMemrefI32(memref<*xi32>) attributes { llvm.emit_c_interface }
 
-!Zp1 = !mod_arith.mod_arith<7681 : i26>
+!Zp1 = !mod_arith.int<7681 : i26>
 !Zp1v = tensor<6x!Zp1>
 // 33554431 = 2 ** 25 - 1
-!Zp2 = !mod_arith.mod_arith<33554431 : i26>
+!Zp2 = !mod_arith.int<33554431 : i26>
 !Zp2v = tensor<6x!Zp2>
 
 func.func @test_lower_reduce() {

diff --git a/tests/Dialect/ModArith/Conversions/mod_arith_to_arith/runner/lower_sub.mlir b/tests/Dialect/ModArith/Conversions/mod_arith_to_arith/runner/lower_sub.mlir
@@ -5,7 +5,7 @@
 
 func.func private @printMemrefI32(memref<*xi32>) attributes { llvm.emit_c_interface }
 
-!Zp = !mod_arith.mod_arith<7681 : i26>
+!Zp = !mod_arith.int<7681 : i26>
 !Zpv = tensor<4x!Zp>
 
 func.func @test_lower_sub() {