Repatriate polynomial dialect

Due to some governance discussions in upstream MLIR [1], we decided it would be appropriate to move the polynomial dialect back into the HEIR tree. We also have some plans to make minor changes to the types, for instance changing the coefficient type to AnyType and removing the coefficient modulus to support things like RNS+modarith coefficient rings, and we felt it would be simpler to just make these changes in tree rather than burden the MLIR community and core maintainers with them. Cf. #1085 [1]: https://discourse.llvm.org/t/rfc-mlir-project-charter-and-restructuring/82896 PiperOrigin-RevId: 696112603
google · Nov 13, 2024 · 5dac435 · 5dac435
1 parent 58862fb
commit 5dac435
Show file tree

Hide file tree

Showing 52 changed files with 1,909 additions and 102 deletions.
diff --git a/lib/Dialect/BGV/Conversions/BGVToLWE/BUILD b/lib/Dialect/BGV/Conversions/BGVToLWE/BUILD
@@ -15,11 +15,8 @@ cc_library(
         ":pass_inc_gen",
         "@heir//lib/Dialect/BGV/IR:Dialect",
         "@heir//lib/Dialect/LWE/IR:Dialect",
-        "@heir//lib/Utils/ConversionUtils",
-        "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
-        "@llvm-project//mlir:PolynomialDialect",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TensorDialect",
         "@llvm-project//mlir:TransformUtils",

diff --git a/lib/Dialect/BGV/IR/BGVOps.td b/lib/Dialect/BGV/IR/BGVOps.td
@@ -8,7 +8,7 @@ include "mlir/Interfaces/InferTypeOpInterface.td"
 
 include "lib/Dialect/LWE/IR/LWETypes.td"
 include "lib/Dialect/LWE/IR/LWETraits.td"
-include "mlir/Dialect/Polynomial/IR/PolynomialAttributes.td"
+include "lib/Dialect/Polynomial/IR/PolynomialAttributes.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 class BGV_Op<string mnemonic, list<Trait> traits = []> :

diff --git a/lib/Dialect/BGV/IR/BUILD b/lib/Dialect/BGV/IR/BUILD
@@ -91,6 +91,6 @@ gentbl_cc_library(
         ":dialect_inc_gen",
         ":td_files",
         "@heir//lib/Dialect/LWE/IR:td_files",
-        "@llvm-project//mlir:PolynomialTdFiles",
+        "@heir//lib/Dialect/Polynomial/IR:td_files",
     ],
 )
diff --git a/lib/Dialect/CGGI/IR/BUILD b/lib/Dialect/CGGI/IR/BUILD
@@ -96,7 +96,7 @@ gentbl_cc_library(
         ":td_files",
         "@heir//lib/Dialect:td_files",
         "@heir//lib/Dialect/LWE/IR:td_files",
-        "@llvm-project//mlir:PolynomialTdFiles",
+        "@heir//lib/Dialect/Polynomial/IR:td_files",
     ],
 )
 

diff --git a/lib/Dialect/CGGI/IR/CGGIOps.td b/lib/Dialect/CGGI/IR/CGGIOps.td
@@ -5,7 +5,7 @@ include "lib/Dialect/HEIRInterfaces.td"
 include "lib/Dialect/CGGI/IR/CGGIDialect.td"
 include "lib/Dialect/CGGI/IR/CGGIAttributes.td"
 
-include "mlir/Dialect/Polynomial/IR/PolynomialAttributes.td"
+include "lib/Dialect/Polynomial/IR/PolynomialAttributes.td"
 include "lib/Dialect/LWE/IR/LWETypes.td"
 
 include "mlir/IR/OpBase.td"

diff --git a/lib/Dialect/CGGI/Transforms/BUILD b/lib/Dialect/CGGI/Transforms/BUILD
@@ -28,10 +28,10 @@ cc_library(
         ":pass_inc_gen",
         "@heir//lib/Dialect/CGGI/IR:Dialect",
         "@heir//lib/Dialect/LWE/IR:Dialect",
+        "@heir//lib/Dialect/Polynomial/IR:Dialect",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
-        "@llvm-project//mlir:PolynomialDialect",
     ],
 )
 

diff --git a/lib/Dialect/CGGI/Transforms/SetDefaultParameters.cpp b/lib/Dialect/CGGI/Transforms/SetDefaultParameters.cpp
@@ -5,10 +5,10 @@
 #include "lib/Dialect/CGGI/IR/CGGIAttributes.h"
 #include "lib/Dialect/CGGI/IR/CGGIOps.h"
 #include "lib/Dialect/LWE/IR/LWEAttributes.h"
-#include "llvm/include/llvm/ADT/TypeSwitch.h"  // from @llvm-project
-#include "llvm/include/llvm/Support/Debug.h"   // from @llvm-project
-#include "mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.h"  // from @llvm-project
-#include "mlir/include/mlir/Dialect/Polynomial/IR/PolynomialAttributes.h"  // from @llvm-project
+#include "lib/Dialect/Polynomial/IR/Polynomial.h"
+#include "lib/Dialect/Polynomial/IR/PolynomialAttributes.h"
+#include "llvm/include/llvm/ADT/TypeSwitch.h"        // from @llvm-project
+#include "llvm/include/llvm/Support/Debug.h"         // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinAttributes.h"  // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinTypes.h"       // from @llvm-project
 #include "mlir/include/mlir/IR/Visitors.h"           // from @llvm-project

diff --git a/lib/Dialect/CKKS/IR/BUILD b/lib/Dialect/CKKS/IR/BUILD
@@ -91,6 +91,6 @@ gentbl_cc_library(
         ":dialect_inc_gen",
         ":td_files",
         "@heir//lib/Dialect/LWE/IR:td_files",
-        "@llvm-project//mlir:PolynomialTdFiles",
+        "@heir//lib/Dialect/Polynomial/IR:td_files",
     ],
 )
diff --git a/lib/Dialect/CKKS/IR/CKKSOps.td b/lib/Dialect/CKKS/IR/CKKSOps.td
@@ -8,7 +8,7 @@ include "mlir/Interfaces/InferTypeOpInterface.td"
 
 include "lib/Dialect/LWE/IR/LWETypes.td"
 include "lib/Dialect/LWE/IR/LWETraits.td"
-include "mlir/Dialect/Polynomial/IR/PolynomialAttributes.td"
+include "lib/Dialect/Polynomial/IR/PolynomialAttributes.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 class CKKS_Op<string mnemonic, list<Trait> traits = []> :

diff --git a/lib/Dialect/LWE/Conversions/LWEToPolynomial/BUILD b/lib/Dialect/LWE/Conversions/LWEToPolynomial/BUILD
@@ -14,13 +14,13 @@ cc_library(
     deps = [
         ":pass_inc_gen",
         "@heir//lib/Dialect/LWE/IR:Dialect",
+        "@heir//lib/Dialect/Polynomial/IR:Dialect",
         "@heir//lib/Dialect/Random/IR:Dialect",
         "@heir//lib/Utils/ConversionUtils",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
-        "@llvm-project//mlir:PolynomialDialect",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TensorDialect",
         "@llvm-project//mlir:TransformUtils",

diff --git a/lib/Dialect/LWE/Conversions/LWEToPolynomial/LWEToPolynomial.cpp b/lib/Dialect/LWE/Conversions/LWEToPolynomial/LWEToPolynomial.cpp
@@ -7,19 +7,19 @@
 #include "lib/Dialect/LWE/IR/LWEAttributes.h"
 #include "lib/Dialect/LWE/IR/LWEOps.h"
 #include "lib/Dialect/LWE/IR/LWETypes.h"
+#include "lib/Dialect/Polynomial/IR/Polynomial.h"
+#include "lib/Dialect/Polynomial/IR/PolynomialAttributes.h"
+#include "lib/Dialect/Polynomial/IR/PolynomialOps.h"
+#include "lib/Dialect/Polynomial/IR/PolynomialTypes.h"
 #include "lib/Dialect/Random/IR/RandomDialect.h"
 #include "lib/Dialect/Random/IR/RandomEnums.h"
 #include "lib/Dialect/Random/IR/RandomOps.h"
 #include "lib/Dialect/Random/IR/RandomTypes.h"
 #include "lib/Utils/ConversionUtils/ConversionUtils.h"
-#include "llvm/include/llvm/ADT/ArrayRef.h"            // from @llvm-project
-#include "llvm/include/llvm/ADT/TypeSwitch.h"          // from @llvm-project
-#include "llvm/include/llvm/Support/ErrorHandling.h"   // from @llvm-project
-#include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
-#include "mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.h"  // from @llvm-project
-#include "mlir/include/mlir/Dialect/Polynomial/IR/PolynomialAttributes.h"  // from @llvm-project
-#include "mlir/include/mlir/Dialect/Polynomial/IR/PolynomialOps.h"  // from @llvm-project
-#include "mlir/include/mlir/Dialect/Polynomial/IR/PolynomialTypes.h"  // from @llvm-project
+#include "llvm/include/llvm/ADT/ArrayRef.h"              // from @llvm-project
+#include "llvm/include/llvm/ADT/TypeSwitch.h"            // from @llvm-project
+#include "llvm/include/llvm/Support/ErrorHandling.h"     // from @llvm-project
+#include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"    // from @llvm-project
 #include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinOps.h"             // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinTypes.h"           // from @llvm-project

diff --git a/lib/Dialect/LWE/IR/BUILD b/lib/Dialect/LWE/IR/BUILD
@@ -22,10 +22,10 @@ cc_library(
         ":dialect_inc_gen",
         ":ops_inc_gen",
         ":types_inc_gen",
+        "@heir//lib/Dialect/Polynomial/IR:Dialect",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:InferTypeOpInterface",
-        "@llvm-project//mlir:PolynomialDialect",
         "@llvm-project//mlir:Support",
     ],
 )
@@ -159,7 +159,7 @@ gentbl_cc_library(
         ":attributes_inc_gen",
         ":dialect_inc_gen",
         ":td_files",
-        "@llvm-project//mlir:PolynomialTdFiles",
+        "@heir//lib/Dialect/Polynomial/IR:td_files",
     ],
 )
 
@@ -185,6 +185,6 @@ gentbl_cc_library(
         ":dialect_inc_gen",
         ":td_files",
         ":types_inc_gen",
-        "@llvm-project//mlir:PolynomialTdFiles",
+        "@heir//lib/Dialect/Polynomial/IR:td_files",
     ],
 )
diff --git a/lib/Dialect/LWE/IR/LWEAttributes.h b/lib/Dialect/LWE/IR/LWEAttributes.h
@@ -5,7 +5,7 @@
 #include "lib/Dialect/LWE/IR/LWEEnums.h.inc"
 #include "mlir/include/mlir/IR/TensorEncoding.h"  // from @llvm-project
 // Required to pull in poly's Ring_Attr
-#include "mlir/include/mlir/Dialect/Polynomial/IR/PolynomialAttributes.h"  // from @llvm-project
+#include "lib/Dialect/Polynomial/IR/PolynomialAttributes.h"
 #define GET_ATTRDEF_CLASSES
 #include "lib/Dialect/LWE/IR/LWEAttributes.h.inc"
 

diff --git a/lib/Dialect/LWE/IR/LWEDialect.cpp b/lib/Dialect/LWE/IR/LWEDialect.cpp
@@ -6,12 +6,12 @@
 #include "lib/Dialect/LWE/IR/LWEAttributes.h"
 #include "lib/Dialect/LWE/IR/LWEOps.h"
 #include "lib/Dialect/LWE/IR/LWETypes.h"
-#include "llvm/include/llvm/ADT/STLFunctionalExtras.h"  // from @llvm-project
-#include "llvm/include/llvm/ADT/TypeSwitch.h"           // from @llvm-project
-#include "llvm/include/llvm/Support/Casting.h"          // from @llvm-project
-#include "llvm/include/llvm/Support/ErrorHandling.h"    // from @llvm-project
-#include "mlir/include/mlir/Dialect/Polynomial/IR/PolynomialAttributes.h"  // from @llvm-project
-#include "mlir/include/mlir/Dialect/Polynomial/IR/PolynomialTypes.h"  // from @llvm-project
+#include "lib/Dialect/Polynomial/IR/PolynomialAttributes.h"
+#include "lib/Dialect/Polynomial/IR/PolynomialTypes.h"
+#include "llvm/include/llvm/ADT/STLFunctionalExtras.h"   // from @llvm-project
+#include "llvm/include/llvm/ADT/TypeSwitch.h"            // from @llvm-project
+#include "llvm/include/llvm/Support/Casting.h"           // from @llvm-project
+#include "llvm/include/llvm/Support/ErrorHandling.h"     // from @llvm-project
 #include "mlir/include/mlir/IR/Diagnostics.h"            // from @llvm-project
 #include "mlir/include/mlir/IR/DialectImplementation.h"  // from @llvm-project
 

diff --git a/lib/Dialect/LWE/IR/LWEOps.td b/lib/Dialect/LWE/IR/LWEOps.td
@@ -4,7 +4,7 @@
 include "lib/Dialect/LWE/IR/LWEDialect.td"
 include "lib/Dialect/LWE/IR/LWETraits.td"
 include "lib/Dialect/LWE/IR/LWETypes.td"
-include "mlir/Dialect/Polynomial/IR/PolynomialAttributes.td"
+include "lib/Dialect/Polynomial/IR/PolynomialAttributes.td"
 
 include "mlir/IR/BuiltinAttributeInterfaces.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"

diff --git a/lib/Dialect/LWE/IR/LWETraits.h b/lib/Dialect/LWE/IR/LWETraits.h
@@ -2,7 +2,7 @@
 #define LIB_DIALECT_LWE_IR_LWETRAITS_H_
 
 #include "lib/Dialect/LWE/IR/LWETypes.h"
-#include "mlir/include/mlir/Dialect/Polynomial/IR/PolynomialAttributes.h"  // from @llvm-project
+#include "lib/Dialect/Polynomial/IR/PolynomialAttributes.h"
 #include "mlir/include/mlir/IR/OpDefinition.h"        // from @llvm-project
 #include "mlir/include/mlir/IR/Operation.h"           // from @llvm-project
 #include "mlir/include/mlir/Support/LLVM.h"           // from @llvm-project

diff --git a/lib/Dialect/ModArith/IR/BUILD b/lib/Dialect/ModArith/IR/BUILD
@@ -88,6 +88,6 @@ gentbl_cc_library(
     deps = [
         ":dialect_inc_gen",
         ":td_files",
-        "@llvm-project//mlir:PolynomialTdFiles",
+        "@heir//lib/Dialect/Polynomial/IR:td_files",
     ],
 )
diff --git a/lib/Dialect/Polynomial/Conversions/PolynomialToStandard/BUILD b/lib/Dialect/Polynomial/Conversions/PolynomialToStandard/BUILD
@@ -13,6 +13,7 @@ cc_library(
     ],
     deps = [
         ":pass_inc_gen",
+        "@heir//lib/Dialect/Polynomial/IR:Dialect",
         "@heir//lib/Utils/ConversionUtils",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:AffineDialect",
@@ -24,7 +25,6 @@ cc_library(
         "@llvm-project//mlir:LLVMDialect",
         "@llvm-project//mlir:LinalgDialect",
         "@llvm-project//mlir:Pass",
-        "@llvm-project//mlir:PolynomialDialect",
         "@llvm-project//mlir:SCFDialect",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TensorDialect",

diff --git a/lib/Dialect/Polynomial/Conversions/PolynomialToStandard/PolynomialToStandard.cpp b/lib/Dialect/Polynomial/Conversions/PolynomialToStandard/PolynomialToStandard.cpp
@@ -9,6 +9,11 @@
 #include <utility>
 #include <vector>
 
+#include "lib/Dialect/Polynomial/IR/Polynomial.h"
+#include "lib/Dialect/Polynomial/IR/PolynomialAttributes.h"
+#include "lib/Dialect/Polynomial/IR/PolynomialDialect.h"
+#include "lib/Dialect/Polynomial/IR/PolynomialOps.h"
+#include "lib/Dialect/Polynomial/IR/PolynomialTypes.h"
 #include "lib/Utils/ConversionUtils/ConversionUtils.h"
 #include "llvm/include/llvm/Support/Casting.h"         // from @llvm-project
 #include "llvm/include/llvm/Support/Debug.h"           // from @llvm-project
@@ -20,13 +25,8 @@
 #include "mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h"    // from @llvm-project
 #include "mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h"  // from @llvm-project
 #include "mlir/include/mlir/Dialect/Linalg/IR/Linalg.h"    // from @llvm-project
-#include "mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.h"  // from @llvm-project
-#include "mlir/include/mlir/Dialect/Polynomial/IR/PolynomialAttributes.h"  // from @llvm-project
-#include "mlir/include/mlir/Dialect/Polynomial/IR/PolynomialDialect.h"  // from @llvm-project
-#include "mlir/include/mlir/Dialect/Polynomial/IR/PolynomialOps.h"  // from @llvm-project
-#include "mlir/include/mlir/Dialect/Polynomial/IR/PolynomialTypes.h"  // from @llvm-project
-#include "mlir/include/mlir/Dialect/SCF/IR/SCF.h"        // from @llvm-project
-#include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
+#include "mlir/include/mlir/Dialect/SCF/IR/SCF.h"          // from @llvm-project
+#include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.h"    // from @llvm-project
 #include "mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h"  // from @llvm-project
 #include "mlir/include/mlir/IR/AffineExpr.h"             // from @llvm-project
 #include "mlir/include/mlir/IR/AffineMap.h"              // from @llvm-project