ModArith: Support RNS type in ops

Author: znssong

lib/Dialect/ModArith/IR/BUILD

@@ -24,6 +24,7 @@ cc_library(
         ":dialect_inc_gen",
         ":ops_inc_gen",
         ":types_inc_gen",
+        "@heir//lib/Dialect/RNS/IR:Dialect",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:CommonFolders",
@@ -44,7 +45,9 @@ td_library(
     # include from the heir-root to enable fully-qualified include-paths
     includes = ["../../../.."],
     deps = [
+        "@heir//lib/Dialect/RNS/IR:td_files",
         "@heir//lib/Utils/DRR",
+        "@llvm-project//mlir:ArithOpsTdFiles",
         "@llvm-project//mlir:BuiltinDialectTdFiles",
         "@llvm-project//mlir:InferTypeOpInterfaceTdFiles",
         "@llvm-project//mlir:OpBaseTdFiles",
@@ -89,6 +92,5 @@ gentbl_cc_library(
     td_file = "ModArithCanonicalization.td",
     deps = [
         ":td_files",
-        "@llvm-project//mlir:ArithOpsTdFiles",
     ],
 )

lib/Dialect/ModArith/IR/ModArithDialect.td

@@ -16,6 +16,7 @@ def ModArith_Dialect : Dialect {
 
   let dependentDialects = [
     "arith::ArithDialect",
+    "mlir::heir::rns::RNSDialect"
   ];
 }
 

lib/Dialect/ModArith/IR/ModArithOps.cpp

@@ -5,6 +5,7 @@
 #include <cstdint>
 #include <vector>
 
+#include "lib/Dialect/RNS/IR/RNSTypes.h"
 #include "llvm/include/llvm/Support/Debug.h"             // from @llvm-project
 #include "mlir/include/mlir/IR/Builders.h"               // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinAttributes.h"      // from @llvm-project
@@ -46,6 +47,30 @@ LogicalResult verifyModArithType(OpType op, ModArithType type) {
   return success();
 }
 
+/// Ensures that the underlying integer type is wide enough for the coefficient
+template <typename OpType>
+LogicalResult verifyRNSType(OpType op, rns::RNSType type) {
+  for (auto basisType : type.getBasisTypes()) {
+    if (auto modArithType = dyn_cast<ModArithType>(basisType)) {
+      if (failed(verifyModArithType(op, modArithType))) {
+        return op.emitOpError()
+               << "Every basis type in the RNS type must be a valid ModArith.";
+      }
+    } else {
+      return op.emitOpError() << "Unsupported RNS type.";
+    }
+  }
+  return success();
+}
+
+template <typename OpType>
+LogicalResult verifyModArithOrRNSType(OpType op, Type type) {
+  if (auto rnsType = dyn_cast<rns::RNSType>(type)) {
+    return verifyRNSType(op, rnsType);
+  }
+  return verifyModArithType(op, cast<ModArithType>(getElementTypeOrSelf(type)));
+}
+
 template <typename OpType>
 LogicalResult verifySameWidth(OpType op, ModArithType modArithType,
                               IntegerType integerType) {
@@ -72,7 +97,7 @@ LogicalResult ReduceOp::verify() {
 }
 
 LogicalResult AddOp::verify() {
-  return verifyModArithType(*this, getResultModArithType(*this));
+  return verifyModArithOrRNSType(*this, this->getResult().getType());
 }
 
 LogicalResult SubOp::verify() {

lib/Dialect/ModArith/IR/ModArithOps.h

@@ -1,6 +1,8 @@
 #ifndef LIB_DIALECT_MODARITH_IR_MODARITHOPS_H_
 #define LIB_DIALECT_MODARITH_IR_MODARITHOPS_H_
 
+#include "lib/Dialect/RNS/IR/RNSTypes.h"
+
 // NOLINTBEGIN(misc-include-cleaner): Required to define ModArithOps
 #include "lib/Dialect/ModArith/IR/ModArithDialect.h"
 #include "lib/Dialect/ModArith/IR/ModArithTypes.h"

lib/Dialect/ModArith/IR/ModArithOps.td

@@ -1,6 +1,7 @@
 #ifndef LIB_DIALECT_MODARITH_IR_MODARITHOPS_TD_
 #define LIB_DIALECT_MODARITH_IR_MODARITHOPS_TD_
 
+include "lib/Dialect/RNS/IR/RNSTypes.td"
 include "lib/Dialect/ModArith/IR/ModArithDialect.td"
 include "lib/Dialect/ModArith/IR/ModArithTypes.td"
 include "mlir/IR/BuiltinAttributes.td"
@@ -9,6 +10,7 @@ include "mlir/IR/OpBase.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
+def ModArithOrRNSLike: TypeOrValueSemanticsContainer<AnyTypeOf<[ModArith_ModArithType, RNS]>, "mod_arith_or_rns-like">;
 
 class ModArith_Op<string mnemonic, list<Trait> traits = [Pure]> :
         Op<ModArith_Dialect, mnemonic, traits> {
@@ -128,8 +130,8 @@ def ModArith_ReduceOp : ModArith_Op<"reduce", [Pure, ElementwiseMappable, SameOp
 
 class ModArith_BinaryOp<string mnemonic, list<Trait> traits = []> :
     ModArith_Op<mnemonic, traits # [SameOperandsAndResultType, Pure, ElementwiseMappable]>,
-    Arguments<(ins ModArithLike:$lhs, ModArithLike:$rhs)>,
-    Results<(outs ModArithLike:$output)> {
+    Arguments<(ins ModArithOrRNSLike:$lhs, ModArithOrRNSLike:$rhs)>,
+    Results<(outs ModArithOrRNSLike:$output)> {
   let hasVerifier = 1;
   let assemblyFormat ="operands attr-dict `:` type($output)";
 }

lib/Dialect/Polynomial/IR/BUILD

@@ -38,6 +38,8 @@ td_library(
     # include from the heir-root to enable fully-qualified include-paths
     includes = ["../../../.."],
     deps = [
+        "@heir//lib/Dialect/ModArith/IR:td_files",
+        "@llvm-project//mlir:ArithOpsTdFiles",
         "@llvm-project//mlir:BuiltinDialectTdFiles",
         "@llvm-project//mlir:InferTypeOpInterfaceTdFiles",
         "@llvm-project//mlir:OpBaseTdFiles",
@@ -82,7 +84,6 @@ add_heir_dialect_library(
     td_file = "PolynomialOps.td",
     deps = [
         ":td_files",
-        "@heir//lib/Dialect/ModArith/IR:td_files",
     ],
 )
 
@@ -93,7 +94,5 @@ gentbl_cc_library(
     td_file = "PolynomialCanonicalization.td",
     deps = [
         ":td_files",
-        "@heir//lib/Dialect/ModArith/IR:td_files",
-        "@llvm-project//mlir:ArithOpsTdFiles",
     ],
 )

tests/Dialect/ModArith/IR/syntax.mlir

@@ -1,7 +1,9 @@
 // RUN: heir-opt --mlir-print-local-scope %s | FileCheck %s
 
 !Zp = !mod_arith.int<17 : i10>
+!Zp2 = !mod_arith.int<257 : i10>
 !Zp_vec = tensor<4x!Zp>
+!rns = !rns.rns<!Zp, !Zp2>
 
 // CHECK: @test_arith_syntax
 func.func @test_arith_syntax() {
@@ -83,3 +85,11 @@ func.func @test_arith_syntax() {
 
   return
 }
+
+// CHECK: @test_rns_syntax
+func.func @test_rns_syntax(%arg0: !rns, %arg1: !rns) -> !rns {
+  // CHECK: mod_arith.add
+  // CHECK-SAME: rns
+  %result = mod_arith.add %arg0, %arg1 : !rns
+  return %result : !rns
+}