Skip to content

Commit a709ba3

Browse files
wtd2copybara-github
wtd2
authored and
copybara-github
committed
fix: implement physical slot replication for bicyclic and tricyclic layout after
matrix multiplication. In the bicyclic layout, we require for each 0 <= slot < numSlot, (ct, slot) is mapped to (slot % n, slot % m). However, after computing through BSGS we can only guarantee for all 0 <= slot < n*m this property preserves (indeed we can guarantee more, but for the tail part it is not correct if n*m does not divide numSlot). In this commit, we add a new relation called `periodic replication relation` that replicate the first copy to all slots periodically. After each matrix multiplication with CRT layouts (bicyclic, tricyclic), we compose this relation to derive the result. The cost of the layout switching will be up to logN rotations. PiperOrigin-RevId: 953525822
1 parent a8488a1 commit a709ba3

28 files changed

Lines changed: 1599 additions & 10 deletions

File tree

‎lib/Kernel/BUILD‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -186,3 +186,57 @@ cc_test(
186186
"@googletest//:gtest_main",
187187
],
188188
)
189+
190+
cc_test(
191+
name = "BicyclicDiagonalMatmulFuzzTest",
192+
srcs = ["BicyclicDiagonalMatmulFuzzTest.cpp"],
193+
deps = [
194+
":AbstractValue",
195+
":ArithmeticDag",
196+
":EvalVisitor",
197+
":KernelImplementation",
198+
"@fuzztest//fuzztest",
199+
"@googletest//:gtest_main",
200+
"@heir//lib/Utils/Layout:Evaluate",
201+
"@heir//lib/Utils/Layout:Utils",
202+
"@llvm-project//mlir:Analysis",
203+
"@llvm-project//mlir:IR",
204+
"@llvm-project//mlir:Support",
205+
],
206+
)
207+
208+
cc_test(
209+
name = "BicyclicMatmulFuzzTest",
210+
srcs = ["BicyclicMatmulFuzzTest.cpp"],
211+
deps = [
212+
":AbstractValue",
213+
":ArithmeticDag",
214+
":EvalVisitor",
215+
":KernelImplementation",
216+
"@fuzztest//fuzztest",
217+
"@googletest//:gtest_main",
218+
"@heir//lib/Utils/Layout:Evaluate",
219+
"@heir//lib/Utils/Layout:Utils",
220+
"@llvm-project//mlir:Analysis",
221+
"@llvm-project//mlir:IR",
222+
"@llvm-project//mlir:Support",
223+
],
224+
)
225+
226+
cc_test(
227+
name = "TricyclicBatchMatmulFuzzTest",
228+
srcs = ["TricyclicBatchMatmulFuzzTest.cpp"],
229+
deps = [
230+
":AbstractValue",
231+
":ArithmeticDag",
232+
":EvalVisitor",
233+
":KernelImplementation",
234+
"@fuzztest//fuzztest",
235+
"@googletest//:gtest_main",
236+
"@heir//lib/Utils/Layout:Evaluate",
237+
"@heir//lib/Utils/Layout:Utils",
238+
"@llvm-project//mlir:Analysis",
239+
"@llvm-project//mlir:IR",
240+
"@llvm-project//mlir:Support",
241+
],
242+
)

‎lib/Kernel/BicyclicDiagonalMatmulFuzzTest.cpp‎

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

‎lib/Kernel/BicyclicMatmulFuzzTest.cpp‎

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