Skip to content

Commit 4aee96c

Browse files
j2kunj2kun
committed
add basic polynomial canonicalization patterns
1 parent 43911fc commit 4aee96c

File tree

6 files changed

+136
-0
lines changed

6 files changed

+136
-0
lines changed

‎mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.td

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -245,6 +245,7 @@ def Polynomial_SubOp : Polynomial_BinaryOp<"sub"> {
245245
%2 = polynomial.sub %0, %1 : !polynomial.polynomial<#ring>
246246
```
247247
}];
248+
let hasCanonicalizer = 1;
248249
}
249250

250251
def Polynomial_MulOp : Polynomial_BinaryOp<"mul", [Commutative]> {
@@ -480,6 +481,7 @@ def Polynomial_NTTOp : Polynomial_Op<"ntt", [Pure]> {
480481
let arguments = (ins Polynomial_PolynomialType:$input);
481482
let results = (outs RankedTensorOf<[AnyInteger]>:$output);
482483
let assemblyFormat = "$input attr-dict `:` qualified(type($input)) `->` type($output)";
484+
let hasCanonicalizer = 1;
483485
let hasVerifier = 1;
484486
}
485487

@@ -498,6 +500,7 @@ def Polynomial_INTTOp : Polynomial_Op<"intt", [Pure]> {
498500
let arguments = (ins RankedTensorOf<[AnyInteger]>:$input);
499501
let results = (outs Polynomial_PolynomialType:$output);
500502
let assemblyFormat = "$input attr-dict `:` qualified(type($input)) `->` type($output)";
503+
let hasCanonicalizer = 1;
501504
let hasVerifier = 1;
502505
}
503506

‎mlir/lib/Dialect/Polynomial/IR/CMakeLists.txt

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,7 @@
1+
set(LLVM_TARGET_DEFINITIONS PolynomialCanonicalization.td)
2+
mlir_tablegen(PolynomialCanonicalization.inc -gen-rewriters)
3+
add_public_tablegen_target(MLIRPolynomialCanonicalizationIncGen)
4+
15
add_mlir_dialect_library(MLIRPolynomialDialect
26
Polynomial.cpp
37
PolynomialAttributes.cpp

‎mlir/lib/Dialect/Polynomial/IR/PolynomialCanonicalization.td

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
//===- PolynomialCanonicalization.td - Polynomial patterns -*- tablegen -*-===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
9+
#ifndef POLYNOMIAL_CANONICALIZATION
10+
#define POLYNOMIAL_CANONICALIZATION
11+
12+
include "mlir/Dialect/Polynomial/IR/Polynomial.td"
13+
include "mlir/Dialect/Arith/IR/ArithOps.td"
14+
include "mlir/IR/OpBase.td"
15+
include "mlir/IR/PatternBase.td"
16+
17+
// TODO: get the proper scalar type from the operand polynomial ring attribute
18+
def SubAsAdd : Pat<
19+
(Polynomial_SubOp $f, $g),
20+
(Polynomial_AddOp $f,
21+
(Polynomial_MulScalarOp $g,
22+
(Arith_ConstantOp
23+
ConstantAttr<I32Attr, "-1">)))>;
24+
25+
def INTTAfterNTT : Pat<
26+
(Polynomial_INTTOp (Polynomial_NTTOp $poly)),
27+
(replaceWithValue $poly),
28+
[]
29+
>;
30+
31+
def NTTAfterINTT : Pat<
32+
(Polynomial_NTTOp (Polynomial_INTTOp $tensor)),
33+
(replaceWithValue $tensor),
34+
[]
35+
>;
36+
37+
#endif // POLYNOMIAL_CANONICALIZATION

‎mlir/lib/Dialect/Polynomial/IR/PolynomialOps.cpp

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,14 @@
77
//===----------------------------------------------------------------------===//
88

99
#include "mlir/Dialect/Polynomial/IR/PolynomialOps.h"
10+
#include "mlir/Dialect/Arith/IR/Arith.h"
1011
#include "mlir/Dialect/Polynomial/IR/Polynomial.h"
1112
#include "mlir/Dialect/Polynomial/IR/PolynomialAttributes.h"
1213
#include "mlir/Dialect/Polynomial/IR/PolynomialTypes.h"
1314
#include "mlir/IR/Builders.h"
1415
#include "mlir/IR/BuiltinTypes.h"
1516
#include "mlir/IR/Dialect.h"
17+
#include "mlir/IR/PatternMatch.h"
1618
#include "mlir/Support/LogicalResult.h"
1719
#include "llvm/ADT/APInt.h"
1820

@@ -183,3 +185,26 @@ LogicalResult INTTOp::verify() {
183185
auto ring = getOutput().getType().getRing();
184186
return verifyNTTOp(this->getOperation(), ring, tensorType);
185187
}
188+
189+
//===----------------------------------------------------------------------===//
190+
// TableGen'd canonicalization patterns
191+
//===----------------------------------------------------------------------===//
192+
193+
namespace {
194+
#include "PolynomialCanonicalization.inc"
195+
} // namespace
196+
197+
void SubOp::getCanonicalizationPatterns(RewritePatternSet &results,
198+
MLIRContext *context) {
199+
populateWithGenerated(results);
200+
}
201+
202+
void NTTOp::getCanonicalizationPatterns(RewritePatternSet &results,
203+
MLIRContext *context) {
204+
populateWithGenerated(results);
205+
}
206+
207+
void INTTOp::getCanonicalizationPatterns(RewritePatternSet &results,
208+
MLIRContext *context) {
209+
populateWithGenerated(results);
210+
}

‎mlir/test/Dialect/Polynomial/canonicalization.mlir

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
// RUN: mlir-opt -canonicalize %s | FileCheck %s
2+
#ntt_poly = #polynomial.int_polynomial<-1 + x**8>
3+
#ntt_ring = #polynomial.ring<coefficientType=i32, coefficientModulus=256, polynomialModulus=#ntt_poly, primitiveRoot=31>
4+
!ntt_poly_ty = !polynomial.polynomial<#ntt_ring>
5+
!tensor_ty = tensor<8xi32, #ntt_ring>
6+
7+
// CHECK-LABEL: @test_canonicalize_intt_after_ntt
8+
// CHECK: (%[[P:.*]]: [[T:.*]]) -> [[T]]
9+
func.func @test_canonicalize_intt_after_ntt(%p0 : !ntt_poly_ty) -> !ntt_poly_ty {
10+
// CHECK-NOT: polynomial.ntt
11+
// CHECK-NOT: polynomial.intt
12+
// CHECK: %[[RESULT:.+]] = polynomial.add %[[P]], %[[P]] : [[T]]
13+
%t0 = polynomial.ntt %p0 : !ntt_poly_ty -> !tensor_ty
14+
%p1 = polynomial.intt %t0: !tensor_ty -> !ntt_poly_ty
15+
%p2 = polynomial.add %p1, %p1 : !ntt_poly_ty
16+
// CHECK: return %[[RESULT]] : [[T]]
17+
return %p2 : !ntt_poly_ty
18+
}
19+
20+
// CHECK-LABEL: @test_canonicalize_ntt_after_intt
21+
// CHECK: (%[[X:.*]]: [[T:.*]]) -> [[T]]
22+
func.func @test_canonicalize_ntt_after_intt(%t0 : !tensor_ty) -> !tensor_ty {
23+
// CHECK-NOT: polynomial.intt
24+
// CHECK-NOT: polynomial.ntt
25+
// CHECK: %[[RESULT:.+]] = arith.addi %[[X]], %[[X]] : [[T]]
26+
%p0 = polynomial.intt %t0 : !tensor_ty -> !ntt_poly_ty
27+
%t1 = polynomial.ntt %p0 : !ntt_poly_ty -> !tensor_ty
28+
%t2 = arith.addi %t1, %t1 : !tensor_ty
29+
// CHECK: return %[[RESULT]] : [[T]]
30+
return %t2 : !tensor_ty
31+
}
32+
33+
#cycl_2048 = #polynomial.int_polynomial<1 + x**1024>
34+
#ring = #polynomial.ring<coefficientType=i32, coefficientModulus=256:i32, polynomialModulus=#cycl_2048>
35+
#one_plus_x_squared = #polynomial.int_polynomial<1 + x**2>
36+
#one_minus_x_squared = #polynomial.int_polynomial<1 + -1x**2>
37+
38+
// CHECK-LABEL: test_canonicalize_sub_power_of_two_cmod
39+
func.func @test_canonicalize_sub_power_of_two_cmod() -> !polynomial.polynomial<#ring> {
40+
%poly0 = polynomial.constant {value=#one_plus_x_squared} : !polynomial.polynomial<#ring>
41+
%poly1 = polynomial.constant {value=#one_minus_x_squared} : !polynomial.polynomial<#ring>
42+
%0 = polynomial.sub %poly0, %poly1 : !polynomial.polynomial<#ring>
43+
// CHECK: %[[minus_one:.+]] = arith.constant -1 : i32
44+
// CHECK: %[[p1:.+]] = polynomial.constant
45+
// CHECK: %[[p2:.+]] = polynomial.constant
46+
// CHECK: %[[p2neg:.+]] = polynomial.mul_scalar %[[p2]], %[[minus_one]]
47+
// CHECK: [[ADD:%.+]] = polynomial.add %[[p1]], %[[p2neg]]
48+
return %0 : !polynomial.polynomial<#ring>
49+
}

‎utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6703,6 +6703,7 @@ cc_library(
67036703
":IR",
67046704
":InferTypeOpInterface",
67056705
":PolynomialAttributesIncGen",
6706+
":PolynomialCanonicalizationIncGen",
67066707
":PolynomialIncGen",
67076708
":Support",
67086709
"//llvm:Support",
@@ -6793,6 +6794,23 @@ gentbl_cc_library(
67936794
deps = [":PolynomialTdFiles"],
67946795
)
67956796

6797+
gentbl_cc_library(
6798+
name = "PolynomialCanonicalizationIncGen",
6799+
strip_include_prefix = "include/mlir/Dialect/Polynomial/IR",
6800+
tbl_outs = [
6801+
(
6802+
["-gen-rewriters"],
6803+
"include/mlir/Dialect/Polynomial/IR/PolynomialCanonicalization.inc",
6804+
),
6805+
],
6806+
tblgen = ":mlir-tblgen",
6807+
td_file = "lib/Dialect/Polynomial/IR/PolynomialCanonicalization.td",
6808+
deps = [
6809+
":ArithOpsTdFiles",
6810+
":PolynomialTdFiles",
6811+
],
6812+
)
6813+
67966814
td_library(
67976815
name = "SPIRVOpsTdFiles",
67986816
srcs = glob(["include/mlir/Dialect/SPIRV/IR/*.td"]),

0 commit comments

Comments
 (0)