1// RUN: mlir-opt %s | FileCheck %s 2 3// This simply tests for syntax. 4 5#my_poly = #polynomial.int_polynomial<1 + x**1024> 6#my_poly_2 = #polynomial.int_polynomial<2> 7#my_poly_3 = #polynomial.int_polynomial<3x> 8#my_poly_4 = #polynomial.int_polynomial<t**3 + 4t + 2> 9#ring1 = #polynomial.ring<coefficientType=i32, coefficientModulus=2837465, polynomialModulus=#my_poly> 10#ring2 = #polynomial.ring<coefficientType=f32> 11#one_plus_x_squared = #polynomial.int_polynomial<1 + x**2> 12 13#ideal = #polynomial.int_polynomial<-1 + x**1024> 14#ring = #polynomial.ring<coefficientType=i32, coefficientModulus=256, polynomialModulus=#ideal> 15!poly_ty = !polynomial.polynomial<ring=#ring> 16 17#ntt_poly = #polynomial.int_polynomial<-1 + x**8> 18#ntt_ring = #polynomial.ring<coefficientType=i32, coefficientModulus=256, polynomialModulus=#ntt_poly> 19!ntt_poly_ty = !polynomial.polynomial<ring=#ntt_ring> 20 21#ntt_poly_2 = #polynomial.int_polynomial<1 + x**65536> 22#ntt_ring_2 = #polynomial.ring<coefficientType = i32, coefficientModulus = 786433 : i32, polynomialModulus=#ntt_poly_2> 23#ntt_ring_2_root = #polynomial.primitive_root<value=283965:i32, degree=131072:i32> 24!ntt_poly_ty_2 = !polynomial.polynomial<ring=#ntt_ring_2> 25 26module { 27 func.func @test_multiply() -> !polynomial.polynomial<ring=#ring1> { 28 %c0 = arith.constant 0 : index 29 %two = arith.constant 2 : i16 30 %five = arith.constant 5 : i16 31 %coeffs1 = tensor.from_elements %two, %two, %five : tensor<3xi16> 32 %coeffs2 = tensor.from_elements %five, %five, %two : tensor<3xi16> 33 34 %poly1 = polynomial.from_tensor %coeffs1 : tensor<3xi16> -> !polynomial.polynomial<ring=#ring1> 35 %poly2 = polynomial.from_tensor %coeffs2 : tensor<3xi16> -> !polynomial.polynomial<ring=#ring1> 36 37 %3 = polynomial.mul %poly1, %poly2 : !polynomial.polynomial<ring=#ring1> 38 39 return %3 : !polynomial.polynomial<ring=#ring1> 40 } 41 42 func.func @test_elementwise(%p0 : !polynomial.polynomial<ring=#ring1>, %p1: !polynomial.polynomial<ring=#ring1>) { 43 %tp0 = tensor.from_elements %p0, %p1 : tensor<2x!polynomial.polynomial<ring=#ring1>> 44 %tp1 = tensor.from_elements %p1, %p0 : tensor<2x!polynomial.polynomial<ring=#ring1>> 45 46 %c = arith.constant 2 : i32 47 %mul_const_sclr = polynomial.mul_scalar %tp0, %c : tensor<2x!polynomial.polynomial<ring=#ring1>>, i32 48 49 %add = polynomial.add %tp0, %tp1 : tensor<2x!polynomial.polynomial<ring=#ring1>> 50 %sub = polynomial.sub %tp0, %tp1 : tensor<2x!polynomial.polynomial<ring=#ring1>> 51 %mul = polynomial.mul %tp0, %tp1 : tensor<2x!polynomial.polynomial<ring=#ring1>> 52 53 return 54 } 55 56 func.func @test_to_from_tensor(%p0 : !polynomial.polynomial<ring=#ring1>) { 57 %c0 = arith.constant 0 : index 58 %two = arith.constant 2 : i16 59 %coeffs1 = tensor.from_elements %two, %two : tensor<2xi16> 60 // CHECK: from_tensor 61 %poly = polynomial.from_tensor %coeffs1 : tensor<2xi16> -> !polynomial.polynomial<ring=#ring1> 62 // CHECK: to_tensor 63 %tensor = polynomial.to_tensor %poly : !polynomial.polynomial<ring=#ring1> -> tensor<1024xi16> 64 65 return 66 } 67 68 func.func @test_degree(%p0 : !polynomial.polynomial<ring=#ring1>) { 69 %0, %1 = polynomial.leading_term %p0 : !polynomial.polynomial<ring=#ring1> -> (index, i32) 70 return 71 } 72 73 func.func @test_monomial() { 74 %deg = arith.constant 1023 : index 75 %five = arith.constant 5 : i16 76 %0 = polynomial.monomial %five, %deg : (i16, index) -> !polynomial.polynomial<ring=#ring1> 77 return 78 } 79 80 func.func @test_monic_monomial_mul() { 81 %five = arith.constant 5 : index 82 %0 = polynomial.constant int<1 + x**2> : !polynomial.polynomial<ring=#ring1> 83 %1 = polynomial.monic_monomial_mul %0, %five : (!polynomial.polynomial<ring=#ring1>, index) -> !polynomial.polynomial<ring=#ring1> 84 return 85 } 86 87 func.func @test_constant() { 88 %0 = polynomial.constant int<1 + x**2> : !polynomial.polynomial<ring=#ring1> 89 %1 = polynomial.constant int<1 + x**2> : !polynomial.polynomial<ring=#ring1> 90 %2 = polynomial.constant float<1.5 + 0.5 x**2> : !polynomial.polynomial<ring=#ring2> 91 92 // Test verbose fallbacks 93 %verb0 = polynomial.constant #polynomial.typed_int_polynomial<1 + x**2> : !polynomial.polynomial<ring=#ring1> 94 %verb2 = polynomial.constant #polynomial.typed_float_polynomial<1.5 + 0.5 x**2> : !polynomial.polynomial<ring=#ring2> 95 return 96 } 97 98 func.func @test_ntt(%0 : !ntt_poly_ty) { 99 %1 = polynomial.ntt %0 {root=#polynomial.primitive_root<value=31:i32, degree=8:index>} : !ntt_poly_ty -> tensor<8xi32, #ntt_ring> 100 return 101 } 102 103 func.func @test_ntt_with_overflowing_root(%0 : !ntt_poly_ty_2) { 104 %1 = polynomial.ntt %0 {root=#ntt_ring_2_root} : !ntt_poly_ty_2 -> tensor<65536xi32, #ntt_ring_2> 105 return 106 } 107 108 func.func @test_intt(%0 : tensor<8xi32, #ntt_ring>) { 109 %1 = polynomial.intt %0 {root=#polynomial.primitive_root<value=31:i32, degree=8:index>} : tensor<8xi32, #ntt_ring> -> !ntt_poly_ty 110 return 111 } 112} 113