Skip to content

Commit 4c27866

Browse files
wtd2copybara-github
wtd2
authored and
copybara-github
committed
feat: add pt-ct matmul support for bicyclic encoding
Define a new kernel MatmulBicyclicDiagonal for pt-ct or ct-pt matrix multiplication. Assign the layout of plaintext operand to generalized diagonal for concrete efficiency: no co-prime property required for plaintext and avoid huge polynomial dimension to store the tensor inside one ciphertext. This can be useful in feed-forward network (FFN) of transformers. Lid tests and fuzz tests added. Fuzz tests added for existing ct-ct matmul (bicyclic) and ct-ct batch matmul (tricyclic). PiperOrigin-RevId: 961125709
1 parent 7bf54eb commit 4c27866

20 files changed

Lines changed: 1304 additions & 6 deletions

File tree

‎lib/Kernel/BUILD‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -200,3 +200,54 @@ cc_test(
200200
"@googletest//:gtest_main",
201201
],
202202
)
203+
204+
cc_test(
205+
name = "BicyclicDiagonalMatmulFuzzTest",
206+
srcs = ["BicyclicDiagonalMatmulFuzzTest.cpp"],
207+
deps = [
208+
":AbstractValue",
209+
":ArithmeticDag",
210+
":EvalVisitor",
211+
":KernelImplementation",
212+
"@fuzztest//fuzztest",
213+
"@googletest//:gtest_main",
214+
"@heir//lib/Utils/Layout:Evaluate",
215+
"@heir//lib/Utils/Layout:Utils",
216+
"@llvm-project//mlir:IR",
217+
"@llvm-project//mlir:Support",
218+
],
219+
)
220+
221+
cc_test(
222+
name = "BicyclicMatmulFuzzTest",
223+
srcs = ["BicyclicMatmulFuzzTest.cpp"],
224+
deps = [
225+
":AbstractValue",
226+
":ArithmeticDag",
227+
":EvalVisitor",
228+
":KernelImplementation",
229+
"@fuzztest//fuzztest",
230+
"@googletest//:gtest_main",
231+
"@heir//lib/Utils/Layout:Evaluate",
232+
"@heir//lib/Utils/Layout:Utils",
233+
"@llvm-project//mlir:IR",
234+
"@llvm-project//mlir:Support",
235+
],
236+
)
237+
238+
cc_test(
239+
name = "TricyclicBatchMatmulFuzzTest",
240+
srcs = ["TricyclicBatchMatmulFuzzTest.cpp"],
241+
deps = [
242+
":AbstractValue",
243+
":ArithmeticDag",
244+
":EvalVisitor",
245+
":KernelImplementation",
246+
"@fuzztest//fuzztest",
247+
"@googletest//:gtest_main",
248+
"@heir//lib/Utils/Layout:Evaluate",
249+
"@heir//lib/Utils/Layout:Utils",
250+
"@llvm-project//mlir:IR",
251+
"@llvm-project//mlir:Support",
252+
],
253+
)

‎lib/Kernel/BicyclicDiagonalMatmulFuzzTest.cpp‎

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
#include <cstdint>
2+
#include <numeric>
3+
#include <optional>
4+
#include <tuple>
5+
#include <utility>
6+
#include <vector>
7+
8+
#include "gtest/gtest.h" // from @googletest
9+
#include "lib/Kernel/AbstractValue.h"
10+
#include "lib/Kernel/ArithmeticDag.h"
11+
#include "lib/Kernel/EvalVisitor.h"
12+
#include "lib/Kernel/KernelImplementation.h"
13+
#include "lib/Utils/Layout/Evaluate.h"
14+
#include "lib/Utils/Layout/Utils.h"
15+
#include "mlir/include/mlir/IR/BuiltinTypes.h" // from @llvm-project
16+
#include "mlir/include/mlir/IR/MLIRContext.h" // from @llvm-project
17+
18+
// copybara hack: avoid reordering include
19+
#include "fuzztest/fuzztest.h" // from @fuzztest
20+
21+
namespace mlir {
22+
namespace heir {
23+
namespace kernel {
24+
namespace {
25+
26+
std::vector<std::vector<int>> runDiagonalMatmul(bool isCtPt,
27+
const std::vector<int>& vecCt,
28+
const std::vector<int>& vecPt,
29+
int64_t m, int64_t n, int64_t p,
30+
bool unroll = true) {
31+
MLIRContext context;
32+
int64_t numSlots = 2 * m * n * p;
33+
34+
int64_t rowsCt = isCtPt ? m : n;
35+
int64_t colsCt = isCtPt ? n : p;
36+
37+
auto layoutCt = getBicyclicLayoutRelation(
38+
RankedTensorType::get({rowsCt, colsCt}, mlir::IndexType::get(&context)),
39+
numSlots);
40+
41+
auto packedCt =
42+
evaluateLayout<int>(layoutCt, [&](const std::vector<int64_t>& pt) {
43+
return vecCt[pt[0] * colsCt + pt[1]];
44+
});
45+
46+
int64_t rowsPt = isCtPt ? n : m;
47+
int64_t colsPt = isCtPt ? p : n;
48+
int64_t contractionDim = isCtPt ? 0 : 1;
49+
int64_t stride = isCtPt ? m : p;
50+
int64_t steps = n;
51+
int64_t period = isCtPt ? m : p;
52+
53+
RankedTensorType matrixType =
54+
RankedTensorType::get({rowsPt, colsPt}, mlir::IndexType::get(&context));
55+
auto layoutPt =
56+
getBicyclicDiagonalRelation(matrixType, contractionDim, stride, numSlots);
57+
auto packedPt =
58+
evaluateLayout<int>(layoutPt, [&](const std::vector<int64_t>& pt) {
59+
return vecPt[pt[0] * colsPt + pt[1]];
60+
});
61+
62+
LiteralValue secretVal = packedCt[0];
63+
LiteralValue plainVal = packedPt;
64+
65+
auto dag = implementRotateAndReduce(
66+
secretVal, std::optional<LiteralValue>(plainVal), period, steps,
67+
DagType::intTensor(32, {numSlots}), {}, "arith.addi", unroll);
68+
69+
LiteralValue result = evalKernel(dag)[0];
70+
auto resultVec = std::get<std::vector<int>>(result.get());
71+
72+
auto resultLayout = getBicyclicLayoutRelation(
73+
RankedTensorType::get({m, p}, mlir::IndexType::get(&context)), numSlots);
74+
75+
return unpackLayoutToMatrix<int>(resultLayout, {resultVec}, {m, p});
76+
}
77+
78+
void diagonalMatmulMatchesNaive(
79+
const std::tuple<bool, int64_t, int64_t, int64_t, std::vector<int>,
80+
std::vector<int>>& args,
81+
bool unroll) {
82+
const auto& [isCtPt, m, n, p, vecCt, vecPt] = args;
83+
84+
if (std::gcd(m, n) != 1 || std::gcd(n, p) != 1 || std::gcd(m, p) != 1) return;
85+
86+
std::vector<std::vector<int>> expected(m, std::vector<int>(p, 0));
87+
if (isCtPt) {
88+
for (int64_t i = 0; i < m; ++i) {
89+
for (int64_t j = 0; j < p; ++j) {
90+
for (int64_t k = 0; k < n; ++k) {
91+
expected[i][j] += vecCt[i * n + k] * vecPt[k * p + j];
92+
}
93+
}
94+
}
95+
} else {
96+
for (int64_t i = 0; i < m; ++i) {
97+
for (int64_t j = 0; j < p; ++j) {
98+
for (int64_t k = 0; k < n; ++k) {
99+
expected[i][j] += vecPt[i * n + k] * vecCt[k * p + j];
100+
}
101+
}
102+
}
103+
}
104+
105+
std::vector<std::vector<int>> actual =
106+
runDiagonalMatmul(isCtPt, vecCt, vecPt, m, n, p, unroll);
107+
108+
EXPECT_EQ(expected, actual);
109+
}
110+
111+
auto shapeAndMatrices() {
112+
return fuzztest::FlatMap(
113+
[](bool isCtPt, int64_t m, int64_t n, int64_t p) {
114+
int64_t sizeCt = isCtPt ? m * n : n * p;
115+
int64_t sizePt = isCtPt ? n * p : m * n;
116+
return fuzztest::TupleOf(
117+
fuzztest::Just(isCtPt), fuzztest::Just(m), fuzztest::Just(n),
118+
fuzztest::Just(p),
119+
/*vecCt=*/
120+
fuzztest::VectorOf(fuzztest::InRange(-100, 100)).WithSize(sizeCt),
121+
/*vecPt=*/
122+
fuzztest::VectorOf(fuzztest::InRange(-100, 100)).WithSize(sizePt));
123+
},
124+
/*isCtPt=*/fuzztest::Arbitrary<bool>(),
125+
/*m=*/fuzztest::InRange<int64_t>(1, 16),
126+
/*n=*/fuzztest::InRange<int64_t>(1, 16),
127+
/*p=*/fuzztest::InRange<int64_t>(1, 16));
128+
}
129+
130+
TEST(BicyclicDiagonalMatmulFuzzTest, CtPtRegression) {
131+
diagonalMatmulMatchesNaive(
132+
{true,
133+
3,
134+
5,
135+
2,
136+
{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15},
137+
{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}},
138+
/*unroll=*/true);
139+
}
140+
141+
TEST(BicyclicDiagonalMatmulFuzzTest, PtCtRegression) {
142+
diagonalMatmulMatchesNaive(
143+
{false,
144+
3,
145+
5,
146+
2,
147+
{1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
148+
{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}},
149+
/*unroll=*/true);
150+
}
151+
152+
TEST(BicyclicDiagonalMatmulFuzzTest, UnitRowDimRegression) {
153+
diagonalMatmulMatchesNaive(
154+
{true, 1, 5, 2, {1, 2, 3, 4, 5}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10}},
155+
/*unroll=*/true);
156+
}
157+
158+
TEST(BicyclicDiagonalMatmulFuzzTest, UnitContractionDimRegression) {
159+
diagonalMatmulMatchesNaive({true, 3, 1, 2, {1, 2, 3}, {4, 5}},
160+
/*unroll=*/true);
161+
}
162+
163+
FUZZ_TEST(BicyclicDiagonalMatmulFuzzTest, diagonalMatmulMatchesNaive)
164+
.WithDomains(shapeAndMatrices(), fuzztest::Arbitrary<bool>());
165+
166+
} // namespace
167+
} // namespace kernel
168+
} // namespace heir
169+
} // namespace mlir

‎lib/Kernel/BicyclicMatmulFuzzTest.cpp‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
#include <cstdint>
2+
#include <numeric>
3+
#include <tuple>
4+
#include <utility>
5+
#include <vector>
6+
7+
#include "gtest/gtest.h" // from @googletest
8+
#include "lib/Kernel/AbstractValue.h"
9+
#include "lib/Kernel/ArithmeticDag.h"
10+
#include "lib/Kernel/EvalVisitor.h"
11+
#include "lib/Kernel/KernelImplementation.h"
12+
#include "lib/Utils/Layout/Evaluate.h"
13+
#include "lib/Utils/Layout/Utils.h"
14+
#include "mlir/include/mlir/IR/BuiltinTypes.h" // from @llvm-project
15+
#include "mlir/include/mlir/IR/MLIRContext.h" // from @llvm-project
16+
#include "mlir/include/mlir/Support/LLVM.h" // from @llvm-project
17+
18+
// copybara hack: avoid reordering include
19+
#include "fuzztest/fuzztest.h" // from @fuzztest
20+
21+
namespace mlir {
22+
namespace heir {
23+
namespace kernel {
24+
namespace {
25+
26+
std::vector<std::vector<int>> runBicyclicMatmul(const std::vector<int>& vecA,
27+
const std::vector<int>& vecB,
28+
int64_t m, int64_t n,
29+
int64_t p) {
30+
MLIRContext context;
31+
int64_t numSlots = m * n * p;
32+
33+
auto layoutA = getBicyclicLayoutRelation(
34+
RankedTensorType::get({m, n}, mlir::IndexType::get(&context)), numSlots);
35+
auto packedA = evaluateLayout<int>(
36+
layoutA,
37+
[&](const std::vector<int64_t>& pt) { return vecA[pt[0] * n + pt[1]]; });
38+
39+
auto layoutB = getBicyclicLayoutRelation(
40+
RankedTensorType::get({n, p}, mlir::IndexType::get(&context)), numSlots);
41+
auto packedB = evaluateLayout<int>(
42+
layoutB,
43+
[&](const std::vector<int64_t>& pt) { return vecB[pt[0] * p + pt[1]]; });
44+
45+
LiteralValue packedAValue = packedA[0];
46+
LiteralValue packedBValue = packedB[0];
47+
48+
auto dag = implementBicyclicMatmul(packedAValue, packedBValue, m, n, p,
49+
DagType::intTensor(32, {numSlots}));
50+
LiteralValue result = evalKernel(dag)[0];
51+
auto resultVec = std::get<std::vector<int>>(result.get());
52+
53+
auto resultLayout = getBicyclicLayoutRelation(
54+
RankedTensorType::get({m, p}, mlir::IndexType::get(&context)), numSlots);
55+
return unpackLayoutToMatrix<int>(resultLayout, {resultVec}, {m, p});
56+
}
57+
58+
void bicyclicMatmulMatchesNaive(
59+
const std::tuple<int64_t, int64_t, int64_t, std::vector<int>,
60+
std::vector<int>>& args) {
61+
const auto& [m, n, p, vecA, vecB] = args;
62+
63+
if (std::gcd(m, n) != 1 || std::gcd(n, p) != 1 || std::gcd(m, p) != 1) return;
64+
65+
std::vector<std::vector<int>> expected(m, std::vector<int>(p, 0));
66+
for (int64_t i = 0; i < m; ++i) {
67+
for (int64_t j = 0; j < p; ++j) {
68+
for (int64_t k = 0; k < n; ++k) {
69+
expected[i][j] += vecA[i * n + k] * vecB[k * p + j];
70+
}
71+
}
72+
}
73+
74+
std::vector<std::vector<int>> actual = runBicyclicMatmul(vecA, vecB, m, n, p);
75+
76+
EXPECT_EQ(expected, actual);
77+
}
78+
79+
auto shapeAndMatrices() {
80+
return fuzztest::FlatMap(
81+
[](int64_t m, int64_t n, int64_t p) {
82+
return fuzztest::TupleOf(
83+
fuzztest::Just(m), fuzztest::Just(n), fuzztest::Just(p),
84+
/*vecA=*/
85+
fuzztest::VectorOf(fuzztest::InRange(-100, 100)).WithSize(m * n),
86+
/*vecB=*/
87+
fuzztest::VectorOf(fuzztest::InRange(-100, 100)).WithSize(n * p));
88+
},
89+
/*m=*/fuzztest::InRange<int64_t>(1, 16),
90+
/*n=*/fuzztest::InRange<int64_t>(1, 16),
91+
/*p=*/fuzztest::InRange<int64_t>(1, 16));
92+
}
93+
94+
TEST(BicyclicMatmulFuzzTest, BicyclicMatmulRegression) {
95+
bicyclicMatmulMatchesNaive(
96+
{3,
97+
5,
98+
2,
99+
{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15},
100+
{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}});
101+
}
102+
103+
FUZZ_TEST(BicyclicMatmulFuzzTest, bicyclicMatmulMatchesNaive)
104+
.WithDomains(shapeAndMatrices());
105+
106+
} // namespace
107+
} // namespace kernel
108+
} // namespace heir
109+
} // namespace mlir

‎lib/Kernel/Kernel.cpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@ std::string kernelNameAsStr(const KernelName& kernelName) {
3434
return "VecmatDiagonal";
3535
case KernelName::MatmulBicyclic:
3636
return "MatmulBicyclic";
37+
case KernelName::MatmulBicyclicDiagonal:
38+
return "MatmulBicyclicDiagonal";
3739
case KernelName::BatchMatmulTricyclic:
3840
return "BatchMatmulTricyclic";
3941
case KernelName::Dot:

‎lib/Kernel/Kernel.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@ struct FieldParser<heir::KernelName> {
2929
if (kernelName == "VecmatDiagonal") return heir::KernelName::VecmatDiagonal;
3030
if (kernelName == "MatmulDiagonal") return heir::KernelName::MatmulDiagonal;
3131
if (kernelName == "MatmulBicyclic") return heir::KernelName::MatmulBicyclic;
32+
if (kernelName == "MatmulBicyclicDiagonal")
33+
return heir::KernelName::MatmulBicyclicDiagonal;
3234
if (kernelName == "BatchMatmulTricyclic")
3335
return heir::KernelName::BatchMatmulTricyclic;
3436
if (kernelName == "Dot") return heir::KernelName::Dot;

0 commit comments

Comments
 (0)