feat(compiler): MANP Analysis of HLFHELinalg.apply_lookup_table (close #175)

This commit is contained in:
Quentin Bourgerie
2021-10-27 14:10:04 +02:00
committed by Andi Drebes
parent ccaf1bff15
commit 5bb58453aa
3 changed files with 24 additions and 2 deletions

View File

@@ -680,6 +680,9 @@ struct MANPAnalysis : public mlir::ForwardDataFlowAnalysis<MANPLatticeValue> {
llvm::dyn_cast<mlir::zamalang::HLFHELinalg::MulEintIntOp>(
op)) {
norm2SqEquiv = getSqMANP(mulEintIntOp, operands);
} else if (llvm::isa<mlir::zamalang::HLFHELinalg::ApplyLookupTableEintOp>(
op)) {
norm2SqEquiv = llvm::APInt{1, 1, false};
}
// Tensor Operators
// ExtractOp

View File

@@ -92,3 +92,23 @@ func @chain_add_eint_int(%e: tensor<8x!HLFHE.eint<2>>) -> tensor<8x!HLFHE.eint<2
%3 = "HLFHELinalg.add_eint_int"(%2, %cst3) : (tensor<8x!HLFHE.eint<2>>, tensor<8xi3>) -> tensor<8x!HLFHE.eint<2>>
return %3 : tensor<8x!HLFHE.eint<2>>
}
// -----
func @apply_lookup_table(%t: tensor<3x3x!HLFHE.eint<2>>) -> tensor<3x3x!HLFHE.eint<3>> {
%lut = std.constant dense<[1,3,5,7]> : tensor<4xi64>
// CHECK: %[[RES:.*]] = "HLFHELinalg.apply_lookup_table"(%[[T:.*]], %[[LUT:.*]]) {MANP = 1 : ui1} : (tensor<3x3x!HLFHE.eint<2>>, tensor<4xi64>) -> tensor<3x3x!HLFHE.eint<3>>
%res = "HLFHELinalg.apply_lookup_table"(%t, %lut) : (tensor<3x3x!HLFHE.eint<2>>, tensor<4xi64>) -> tensor<3x3x!HLFHE.eint<3>>
return %res : tensor<3x3x!HLFHE.eint<3>>
}
// -----
func @apply_lookup_table_after_op(%t: tensor<8x!HLFHE.eint<2>>, %i: tensor<8xi3>) -> tensor<8x!HLFHE.eint<3>> {
%lut = std.constant dense<[1,3,5,7]> : tensor<4xi64>
// CHECK: %[[V0:.*]] = "HLFHELinalg.mul_eint_int"([[T:.*]], %[[I:.*]]) {MANP = 8 : ui{{[0-9]+}}} : (tensor<8x!HLFHE.eint<2>>, tensor<8xi3>) -> tensor<8x!HLFHE.eint<2>>
%0 = "HLFHELinalg.mul_eint_int"(%t, %i) : (tensor<8x!HLFHE.eint<2>>, tensor<8xi3>) -> tensor<8x!HLFHE.eint<2>>
// CHECK-NEXT: %[[RES:.*]] = "HLFHELinalg.apply_lookup_table"(%[[V0:.*]], %[[LUT:.*]]) {MANP = 1 : ui1} : (tensor<8x!HLFHE.eint<2>>, tensor<4xi64>) -> tensor<8x!HLFHE.eint<3>>
%res = "HLFHELinalg.apply_lookup_table"(%0, %lut) : (tensor<8x!HLFHE.eint<2>>, tensor<4xi64>) -> tensor<8x!HLFHE.eint<3>>
return %res : tensor<8x!HLFHE.eint<3>>
}

View File

@@ -885,8 +885,7 @@ TEST(End2EndJit_HLFHELinalg, apply_lookup_table) {
%res = "HLFHELinalg.apply_lookup_table"(%t, %lut) : (tensor<3x3x!HLFHE.eint<2>>, tensor<4xi64>) -> tensor<3x3x!HLFHE.eint<3>>
return %res : tensor<3x3x!HLFHE.eint<3>>
}
)XXX",
"main", true);
)XXX");
const uint8_t t[3][3]{
{0, 1, 2},
{3, 0, 1},