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