Add guards to prevent lowering eval-form polynomial constants

Author: crockeea

lib/Dialect/Polynomial/Conversions/PolynomialToModArith/PolynomialToModArith.cpp

@@ -298,6 +298,11 @@ struct ConvertConstant : public OpConversionPattern<ConstantOp> {
           op, "failed to construct common conversion info");
 
     auto typeInfo = res.value();
+    // TODO(#97): support compile-time NTT
+    if (typeInfo.polynomialType.getForm() == Form::EVAL) {
+      return rewriter.notifyMatchFailure(
+          op, "unsupported eval-form polynomial constant");
+    }
 
     auto attr = dyn_cast<TypedIntPolynomialAttr>(op.getValue());
     if (!attr)
@@ -393,6 +398,45 @@ struct ConvertMonomial : public OpConversionPattern<MonomialOp> {
     auto storageTensorType =
         RankedTensorType::get(storageShape, typeInfo.coefficientStorageType);
 
+    // TODO(#97): support compile-time NTT
+    // We don't have proper support for EVAL-form constants, but we can
+    // at least support degree-zero polynomial constants in EVAL form. The
+    // NTT of a degree-zero polynomial is a vector where each coefficient is the
+    // constant term.
+    if (typeInfo.polynomialType.getForm() == Form::EVAL) {
+      auto degree = adaptor.getDegree().getDefiningOp<arith::ConstantIndexOp>();
+      if (!degree || degree.value() != 0) {
+        return rewriter.notifyMatchFailure(
+            op, "unsupported eval-form non-constant monomial");
+      }
+
+      Value result;
+      if (auto modQTy = dyn_cast<ModQTypeInterface>(typeInfo.coefficientType)) {
+        Type extractedType = modQTy.getLoweringType();
+        Value extracted = mod_arith::ExtractOp::create(
+            b, extractedType, adaptor.getCoefficient());
+        Value replicatedStorage;
+        if (isa<ShapedType>(extractedType)) {
+          auto init = tensor::EmptyOp::create(b, storageTensorType.getShape(),
+                                              typeInfo.coefficientStorageType);
+          replicatedStorage = linalg::BroadcastOp::create(b, extracted, init,
+                                                          ArrayRef<int64_t>{0})
+                                  .getResult()[0];
+        } else {
+          replicatedStorage =
+              tensor::SplatOp::create(b, extracted, storageTensorType);
+        }
+        result = mod_arith::EncapsulateOp::create(b, typeInfo.tensorType,
+                                                  replicatedStorage)
+                     .getResult();
+      } else {
+        result = tensor::SplatOp::create(b, adaptor.getCoefficient(),
+                                         typeInfo.tensorType);
+      }
+      rewriter.replaceOp(op, result);
+      return success();
+    }
+
     auto tensor = arith::ConstantOp::create(
         b, DenseElementsAttr::get(
                storageTensorType,

tests/Dialect/Polynomial/Conversions/polynomial_to_mod_arith/eval_form_errors.mlir

@@ -0,0 +1,24 @@
+// RUN: heir-opt --polynomial-to-mod-arith --verify-diagnostics --split-input-file %s
+
+#poly = #polynomial.int_polynomial<1 + x**4>
+#ring = #polynomial.ring<coefficientType=i32, polynomialModulus=#poly>
+!poly_ty = !polynomial.polynomial<ring=#ring, form=eval>
+
+func.func @eval_constant() -> !poly_ty {
+  // expected-error@+1 {{failed to legalize operation}}
+  %0 = polynomial.constant int<1> : !poly_ty
+  return %0 : !poly_ty
+}
+
+// -----
+
+#poly = #polynomial.int_polynomial<1 + x**4>
+#ring = #polynomial.ring<coefficientType=i32, polynomialModulus=#poly>
+!poly_ty = !polynomial.polynomial<ring=#ring, form=eval>
+
+func.func @eval_nonconstant_monomial(%coeff: i32) -> !poly_ty {
+  %c1 = arith.constant 1 : index
+  // expected-error@+1 {{failed to legalize operation}}
+  %0 = polynomial.monomial %coeff, %c1 : (i32, index) -> !poly_ty
+  return %0 : !poly_ty
+}

tests/Dialect/Polynomial/Conversions/polynomial_to_mod_arith/eval_monomial.mlir

@@ -0,0 +1,31 @@
+// RUN: heir-opt --polynomial-to-mod-arith --mlir-print-local-scope %s | FileCheck %s
+
+#poly = #polynomial.int_polynomial<1 + x**4>
+#ring_i32 = #polynomial.ring<coefficientType=i32, polynomialModulus=#poly>
+!poly_i32 = !polynomial.polynomial<ring=#ring_i32, form=eval>
+
+// CHECK: func @eval_constant_monomial_int
+// CHECK-SAME: (%[[ARG0:.*]]: i32)
+func.func @eval_constant_monomial_int(%coeff: i32) -> !poly_i32 {
+  %c0 = arith.constant 0 : index
+  // CHECK: %[[SPLAT:.*]] = tensor.splat %[[ARG0]] : tensor<4xi32>
+  // CHECK: return %[[SPLAT]]
+  %0 = polynomial.monomial %coeff, %c0 : (i32, index) -> !poly_i32
+  return %0 : !poly_i32
+}
+
+!coeff_ty = !mod_arith.int<17 : i32>
+#ring_mod = #polynomial.ring<coefficientType=!coeff_ty, polynomialModulus=#poly>
+!poly_mod = !polynomial.polynomial<ring=#ring_mod, form=eval>
+
+// CHECK: func @eval_constant_monomial_mod
+// CHECK-SAME: (%[[ARG0:.*]]: !mod_arith.int<17 : i32>)
+func.func @eval_constant_monomial_mod(%coeff: !coeff_ty) -> !poly_mod {
+  %c0 = arith.constant 0 : index
+  // CHECK: %[[EXTRACTED:.*]] = mod_arith.extract %[[ARG0]] : !mod_arith.int<17 : i32> -> i32
+  // CHECK: %[[SPLAT:.*]] = tensor.splat %[[EXTRACTED]] : tensor<4xi32>
+  // CHECK: %[[ENCAPSULATED:.*]] = mod_arith.encapsulate %[[SPLAT]] : tensor<4xi32> -> tensor<4x!mod_arith.int<17 : i32>>
+  // CHECK: return %[[ENCAPSULATED]]
+  %0 = polynomial.monomial %coeff, %c0 : (!coeff_ty, index) -> !poly_mod
+  return %0 : !poly_mod
+}