xref: /llvm-project/mlir/test/Dialect/Polynomial/ops_errors.mlir (revision 692ae5443b1778e138527ef55d799a4b535a36f9)
1145176dcSJeremy Kun// RUN: mlir-opt --split-input-file --verify-diagnostics %s
2145176dcSJeremy Kun
32ff43ce8SJeremy Kun#my_poly = #polynomial.int_polynomial<1 + x**1024>
4*692ae544SJeremy Kun#ring = #polynomial.ring<coefficientType=i16>
52ff43ce8SJeremy Kun!ty = !polynomial.polynomial<ring=#ring>
6145176dcSJeremy Kun
7145176dcSJeremy Kunfunc.func @test_from_tensor_too_large_coeffs() {
8145176dcSJeremy Kun  %two = arith.constant 2 : i32
9145176dcSJeremy Kun  %coeffs1 = tensor.from_elements %two, %two : tensor<2xi32>
10145176dcSJeremy Kun  // expected-error@below {{is too large to fit in the coefficients}}
11145176dcSJeremy Kun  // expected-note@below {{rescaled to fit}}
12145176dcSJeremy Kun  %poly = polynomial.from_tensor %coeffs1 : tensor<2xi32> -> !ty
13145176dcSJeremy Kun  return
14145176dcSJeremy Kun}
15145176dcSJeremy Kun
16145176dcSJeremy Kun// -----
17145176dcSJeremy Kun
182ff43ce8SJeremy Kun#my_poly = #polynomial.int_polynomial<1 + x**4>
192ff43ce8SJeremy Kun#ring = #polynomial.ring<coefficientType=i32, coefficientModulus=256:i32, polynomialModulus=#my_poly>
202ff43ce8SJeremy Kun!ty = !polynomial.polynomial<ring=#ring>
21145176dcSJeremy Kunfunc.func @test_from_tensor_wrong_tensor_type() {
22145176dcSJeremy Kun  %two = arith.constant 2 : i32
23145176dcSJeremy Kun  %coeffs1 = tensor.from_elements %two, %two, %two, %two, %two : tensor<5xi32>
242ff43ce8SJeremy Kun  // expected-error@below {{input type 'tensor<5xi32>' does not match output type '!polynomial.polynomial<ring = <coefficientType = i32, coefficientModulus = 256 : i32, polynomialModulus = <1 + x**4>>>'}}
25145176dcSJeremy Kun  // expected-note@below {{at most the degree of the polynomialModulus of the output type's ring attribute}}
26145176dcSJeremy Kun  %poly = polynomial.from_tensor %coeffs1 : tensor<5xi32> -> !ty
27145176dcSJeremy Kun  return
28145176dcSJeremy Kun}
29145176dcSJeremy Kun
30145176dcSJeremy Kun// -----
31145176dcSJeremy Kun
322ff43ce8SJeremy Kun#my_poly = #polynomial.int_polynomial<1 + x**4>
332ff43ce8SJeremy Kun#ring = #polynomial.ring<coefficientType=i32, coefficientModulus=256:i32, polynomialModulus=#my_poly>
342ff43ce8SJeremy Kun!ty = !polynomial.polynomial<ring=#ring>
35145176dcSJeremy Kunfunc.func @test_to_tensor_wrong_output_tensor_type(%arg0 : !ty) {
362ff43ce8SJeremy Kun  // expected-error@below {{input type '!polynomial.polynomial<ring = <coefficientType = i32, coefficientModulus = 256 : i32, polynomialModulus = <1 + x**4>>>' does not match output type 'tensor<5xi32>'}}
37145176dcSJeremy Kun  // expected-note@below {{at most the degree of the polynomialModulus of the input type's ring attribute}}
38145176dcSJeremy Kun  %tensor = polynomial.to_tensor %arg0 : !ty -> tensor<5xi32>
39145176dcSJeremy Kun  return
40145176dcSJeremy Kun}
41145176dcSJeremy Kun
42145176dcSJeremy Kun// -----
43145176dcSJeremy Kun
442ff43ce8SJeremy Kun#my_poly = #polynomial.int_polynomial<1 + x**1024>
452ff43ce8SJeremy Kun#ring = #polynomial.ring<coefficientType=i16, coefficientModulus=256:i32, polynomialModulus=#my_poly>
462ff43ce8SJeremy Kun!ty = !polynomial.polynomial<ring=#ring>
47145176dcSJeremy Kun
48145176dcSJeremy Kunfunc.func @test_mul_scalar_wrong_type(%arg0: !ty) -> !ty {
49145176dcSJeremy Kun  %scalar = arith.constant 2 : i32  // should be i16
50145176dcSJeremy Kun  // expected-error@below {{polynomial coefficient type 'i16' does not match scalar type 'i32'}}
51145176dcSJeremy Kun  %poly = polynomial.mul_scalar %arg0, %scalar : !ty, i32
52145176dcSJeremy Kun  return %poly : !ty
53145176dcSJeremy Kun}
54624c9fc8SJeremy Kun
55624c9fc8SJeremy Kun// -----
56624c9fc8SJeremy Kun
572ff43ce8SJeremy Kun#my_poly = #polynomial.int_polynomial<-1 + x**1024>
581f46729aSJeremy Kun#ring = #polynomial.ring<coefficientType=i16, coefficientModulus=256:i16, polynomialModulus=#my_poly>
592ff43ce8SJeremy Kun!poly_ty = !polynomial.polynomial<ring=#ring>
60624c9fc8SJeremy Kun
61624c9fc8SJeremy Kun// CHECK-NOT: @test_invalid_ntt
62624c9fc8SJeremy Kun// CHECK-NOT: polynomial.ntt
63624c9fc8SJeremy Kunfunc.func @test_invalid_ntt(%0 : !poly_ty) {
64624c9fc8SJeremy Kun  // expected-error@below {{expects a ring encoding to be provided to the tensor}}
651f46729aSJeremy Kun  %1 = polynomial.ntt %0 {root=#polynomial.primitive_root<value=31:i32, degree=8:index>} : !poly_ty -> tensor<1024xi32>
66624c9fc8SJeremy Kun  return
67624c9fc8SJeremy Kun}
68624c9fc8SJeremy Kun
69624c9fc8SJeremy Kun// -----
70624c9fc8SJeremy Kun
712ff43ce8SJeremy Kun#my_poly = #polynomial.int_polynomial<-1 + x**1024>
722ff43ce8SJeremy Kun#ring = #polynomial.ring<coefficientType=i16, coefficientModulus=256:i16, polynomialModulus=#my_poly>
732ff43ce8SJeremy Kun!poly_ty = !polynomial.polynomial<ring=#ring>
74624c9fc8SJeremy Kun
751f46729aSJeremy Kun// CHECK-NOT: @test_invalid_ntt
761f46729aSJeremy Kun// CHECK-NOT: polynomial.ntt
771f46729aSJeremy Kunfunc.func @test_invalid_ntt(%0 : !poly_ty) {
781f46729aSJeremy Kun  // expected-error@below {{tensor encoding is not a ring attribute}}
791f46729aSJeremy Kun  %1 = polynomial.ntt %0 {root=#polynomial.primitive_root<value=31:i32, degree=8:index>} : !poly_ty -> tensor<1024xi32, #my_poly>
80624c9fc8SJeremy Kun  return
81624c9fc8SJeremy Kun}
82624c9fc8SJeremy Kun
83624c9fc8SJeremy Kun// -----
84624c9fc8SJeremy Kun
852ff43ce8SJeremy Kun#my_poly = #polynomial.int_polynomial<-1 + x**1024>
861f46729aSJeremy Kun#ring = #polynomial.ring<coefficientType=i16, coefficientModulus=256:i16, polynomialModulus=#my_poly>
871f46729aSJeremy Kun#ring1 = #polynomial.ring<coefficientType=i16, coefficientModulus=257:i16, polynomialModulus=#my_poly>
881f46729aSJeremy Kun!poly_ty = !polynomial.polynomial<ring=#ring>
891f46729aSJeremy Kun
901f46729aSJeremy Kun// CHECK-NOT: @test_invalid_intt
911f46729aSJeremy Kun// CHECK-NOT: polynomial.intt
921f46729aSJeremy Kunfunc.func @test_invalid_intt(%0 : tensor<1024xi32, #ring1>) {
931f46729aSJeremy Kun  // expected-error@below {{not equivalent to the polynomial ring}}
941f46729aSJeremy Kun  %1 = polynomial.intt %0 {root=#polynomial.primitive_root<value=31:i32, degree=8:index>} : tensor<1024xi32, #ring1> -> !poly_ty
951f46729aSJeremy Kun  return
961f46729aSJeremy Kun}
971f46729aSJeremy Kun
981f46729aSJeremy Kun// -----
991f46729aSJeremy Kun
1001f46729aSJeremy Kun#my_poly = #polynomial.int_polynomial<-1 + x**1024>
1011f46729aSJeremy Kun#ring = #polynomial.ring<coefficientType=i16, coefficientModulus=256:i16, polynomialModulus=#my_poly>
1022ff43ce8SJeremy Kun!poly_ty = !polynomial.polynomial<ring=#ring>
103624c9fc8SJeremy Kun
104624c9fc8SJeremy Kun// CHECK-NOT: @test_invalid_intt
105624c9fc8SJeremy Kun// CHECK-NOT: polynomial.intt
106624c9fc8SJeremy Kunfunc.func @test_invalid_intt(%0 : tensor<1025xi32, #ring>) {
107624c9fc8SJeremy Kun  // expected-error@below {{does not match output type}}
108624c9fc8SJeremy Kun  // expected-note@below {{exactly the degree of the polynomialModulus of the polynomial type's ring attribute}}
1091f46729aSJeremy Kun  %1 = polynomial.intt %0 {root=#polynomial.primitive_root<value=31:i32, degree=8:index>} : tensor<1025xi32, #ring> -> !poly_ty
110624c9fc8SJeremy Kun  return
111624c9fc8SJeremy Kun}
112624c9fc8SJeremy Kun
113624c9fc8SJeremy Kun// -----
114624c9fc8SJeremy Kun
1152ff43ce8SJeremy Kun#my_poly = #polynomial.int_polynomial<-1 + x**8>
116624c9fc8SJeremy Kun// A valid root is 31
1171f46729aSJeremy Kun#ring = #polynomial.ring<coefficientType=i16, coefficientModulus=256:i16, polynomialModulus=#my_poly>
1182ff43ce8SJeremy Kun!poly_ty = !polynomial.polynomial<ring=#ring>
119624c9fc8SJeremy Kun
120624c9fc8SJeremy Kun// CHECK-NOT: @test_invalid_intt
121624c9fc8SJeremy Kun// CHECK-NOT: polynomial.intt
122624c9fc8SJeremy Kunfunc.func @test_invalid_intt(%0 : tensor<8xi32, #ring>) {
1231f46729aSJeremy Kun  // expected-error@below {{provided root 32 is not a primitive root of unity mod 256, with the specified degree 8}}
1241f46729aSJeremy Kun  %1 = polynomial.intt %0 {root=#polynomial.primitive_root<value=32:i16, degree=8:index>} : tensor<8xi32, #ring> -> !poly_ty
125624c9fc8SJeremy Kun  return
126624c9fc8SJeremy Kun}
127