Skip to content

Commit 980e966

Browse files
copybara-githubcopybara-github
committed
Merge pull request #2884 from crockeea:bugfix2
PiperOrigin-RevId: 901339018
2 parents 7423251 + 89d2722 commit 980e966

4 files changed

Lines changed: 61 additions & 12 deletions

File tree

‎lib/Dialect/Polynomial/Transforms/BUILD‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ cc_library(
5858
":pass_inc_gen",
5959
"@com_google_ortools//ortools/sat:cp_model",
6060
"@com_google_ortools//ortools/sat:cp_model_solver",
61+
"@heir//lib/Dialect/ModArith/IR:ModArithTypeInterfaces",
6162
"@heir//lib/Dialect/Polynomial/IR:Dialect",
6263
"@llvm-project//llvm:Support",
6364
"@llvm-project//mlir:FuncDialect",

‎lib/Dialect/Polynomial/Transforms/NTTSolver.cpp‎

Lines changed: 24 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,16 @@
11
#include "lib/Dialect/Polynomial/Transforms/NTTSolver.h"
22

3+
#include <cassert>
4+
5+
#include "lib/Dialect/ModArith/IR/ModArithTypeInterfaces.h"
36
#include "lib/Dialect/Polynomial/IR/PolynomialAttributes.h"
47
#include "lib/Dialect/Polynomial/IR/PolynomialTypes.h"
5-
#include "mlir/include/mlir/IR/BuiltinTypes.h" // from @llvm-project
6-
#include "mlir/include/mlir/IR/Types.h" // from @llvm-project
7-
#include "mlir/include/mlir/IR/Value.h" // from @llvm-project
8-
#include "mlir/include/mlir/Support/LLVM.h" // from @llvm-project
9-
#include "ortools/sat/cp_model.h" // from @com_google_ortools
10-
#include "ortools/sat/cp_model_solver.h" // from @com_google_ortools
8+
#include "mlir/include/mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project
9+
#include "mlir/include/mlir/IR/Types.h" // from @llvm-project
10+
#include "mlir/include/mlir/IR/Value.h" // from @llvm-project
11+
#include "mlir/include/mlir/Support/LLVM.h" // from @llvm-project
12+
#include "ortools/sat/cp_model.h" // from @com_google_ortools
13+
#include "ortools/sat/cp_model_solver.h" // from @com_google_ortools
1114

1215
namespace mlir {
1316
namespace heir {
@@ -42,13 +45,22 @@ bool CPSATSolution::isValid() const {
4245

4346
int getConversionCost(const Value& v) {
4447
Type t = v.getType();
45-
if (auto p = dyn_cast<PolynomialType>(t)) {
46-
return 1;
47-
}
48-
if (auto rt = dyn_cast<RankedTensorType>(t)) {
49-
return rt.getNumElements();
48+
int multiplicity = 1;
49+
50+
if (auto shaped = dyn_cast<ShapedType>(t)) {
51+
multiplicity = shaped.getNumElements();
52+
t = shaped.getElementType();
5053
}
51-
return 0;
54+
55+
auto polyTy = dyn_cast<PolynomialType>(t);
56+
if (!polyTy) return 0;
57+
58+
auto coeffTy = dyn_cast<mod_arith::ModQTypeInterface>(
59+
polyTy.getRing().getCoefficientType());
60+
assert(coeffTy &&
61+
"polynomial coefficient type must implement ModQTypeInterface");
62+
63+
return coeffTy.getNumResidues() * multiplicity;
5264
}
5365

5466
NTTSolver::RepVars& NTTSolver::getOrCreateVars(const Value& v) {

‎lib/Dialect/Polynomial/Transforms/Passes.td‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,9 @@ def PolyMulToNTT : Pass <"convert-polynomial-mul-to-ntt", "func::FuncOp"> {
1515
let dependentDialects = ["mlir::heir::polynomial::PolynomialDialect"];
1616

1717
let statistics = [
18+
// Note that these statistics count the number of ops inserted, *not* the number
19+
// of Zq NTTs/INTTs. These may be different, e.g., if an NTT op takes an RNS
20+
// polynomial.
1821
Statistic<"numNttsInserted", "num-ntts-inserted", "Number of NTT ops inserted">,
1922
Statistic<"numInttsInserted", "num-intts-inserted", "Number of INTT ops inserted">
2023
];

‎tests/Dialect/Polynomial/Transforms/poly_mul_to_ntt_optimal_placement.mlir‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,21 @@
22

33
!Zq0 = !mod_arith.int<1095233372161 : i64>
44
!Zq1 = !mod_arith.int<1032955396097 : i64>
5+
!Zq2 = !mod_arith.int<1002198353921 : i64>
56
#ring_1 = #polynomial.ring<coefficientType = !rns.rns<!Zq0>, polynomialModulus = <1 + x**1024>>
67
!poly_ty_1 = !polynomial.polynomial<ring=#ring_1, form=coeff>
8+
!ntt_poly_ty_1 = !polynomial.polynomial<ring=#ring_1, form=eval>
79

810
#ring_2 = #polynomial.ring<coefficientType = !rns.rns<!Zq0, !Zq1>, polynomialModulus = <1 + x**1024>>
911
!poly_ty_2 = !polynomial.polynomial<ring=#ring_2, form=coeff>
1012

13+
#ring_3 = #polynomial.ring<coefficientType = !rns.rns<!Zq0, !Zq1, !Zq2>, polynomialModulus = <1 + x**1024>>
14+
!poly_ty_3 = !polynomial.polynomial<ring=#ring_3, form=coeff>
15+
16+
#ring_4 = #polynomial.ring<coefficientType = !rns.rns<!Zq1>, polynomialModulus = <1 + x**1024>>
17+
!poly_ty_4 = !polynomial.polynomial<ring=#ring_4, form=coeff>
18+
!ntt_poly_ty_4 = !polynomial.polynomial<ring=#ring_4, form=eval>
19+
1120
module {
1221
// Covers: a shared-dataflow circuit where optimal placement must avoid
1322
// redundant conversions on mul chains while still satisfying coeff-only
@@ -42,4 +51,28 @@ module {
4251
%out = polynomial.convert_basis %prod {targetBasis = !rns.rns<!Zq0>} : !poly_ty_2 -> !poly_ty_1
4352
return %out, %x_t, %shift_t : !poly_ty_1, tensor<1024x!rns.rns<!Zq0, !Zq1>>, tensor<1024x!rns.rns<!Zq0, !Zq1>>
4453
}
54+
55+
// Covers: limb-aware conversion costs can move NTTs past extract_slice.
56+
// Old unweighted costs would hoist one NTT on %x; the weighted model keeps
57+
// %x in coeff form and inserts two cheaper one-limb NTTs on the slices.
58+
// CHECK: func.func @weighted_extract_slice_ntt_placement([[x:%.+]]: [[poly_ty_3:![^ ]+]]) -> ([[ntt_poly_ty_1:![^ ]+]], [[ntt_poly_ty_4:![^ ]+]], tensor<1024x[[RNS3:![^ ]+]]>) {
59+
// CHECK-NOT: polynomial.ntt
60+
// CHECK: [[a:%.+]] = polynomial.extract_slice [[x]] {size = 1 : index, start = 0 : index} : [[poly_ty_3]] -> [[poly_ty_1:![^ ]+]]
61+
// CHECK: [[a_ntt:%.+]] = polynomial.ntt [[a]] : [[poly_ty_1]]
62+
// CHECK: [[b:%.+]] = polynomial.extract_slice [[x]] {size = 1 : index, start = 1 : index} : [[poly_ty_3]] -> [[poly_ty_4:![^ ]+]]
63+
// CHECK: [[b_ntt:%.+]] = polynomial.ntt [[b]] : [[poly_ty_4]]
64+
// CHECK: [[ma:%.+]] = polynomial.mul [[a_ntt]], [[a_ntt]] : [[ntt_poly_ty_1]]
65+
// CHECK: [[mb:%.+]] = polynomial.mul [[b_ntt]], [[b_ntt]] : [[ntt_poly_ty_4]]
66+
// CHECK: [[x_t:%.+]] = polynomial.to_tensor [[x]] : [[poly_ty_3]] -> tensor<1024x[[RNS3]]>
67+
// CHECK: return [[ma]], [[mb]], [[x_t]] : [[ntt_poly_ty_1]], [[ntt_poly_ty_4]], tensor<1024x[[RNS3]]>
68+
func.func @weighted_extract_slice_ntt_placement(%x: !poly_ty_3) -> (!poly_ty_1, !poly_ty_4, tensor<1024x!rns.rns<!Zq0, !Zq1, !Zq2>>) {
69+
%a = polynomial.extract_slice %x {start = 0 : index, size = 1 : index}
70+
: !poly_ty_3 -> !poly_ty_1
71+
%b = polynomial.extract_slice %x {start = 1 : index, size = 1 : index}
72+
: !poly_ty_3 -> !poly_ty_4
73+
%ma = polynomial.mul %a, %a : !poly_ty_1
74+
%mb = polynomial.mul %b, %b : !poly_ty_4
75+
%x_t = polynomial.to_tensor %x : !poly_ty_3 -> tensor<1024x!rns.rns<!Zq0, !Zq1, !Zq2>>
76+
return %ma, %mb, %x_t : !poly_ty_1, !poly_ty_4, tensor<1024x!rns.rns<!Zq0, !Zq1, !Zq2>>
77+
}
4578
}

0 commit comments

Comments
 (0)