Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support for tiling of more FHELinalg operations #779

Merged
merged 8 commits into from
Apr 12, 2024
Merged

Add support for tiling of more FHELinalg operations #779

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
a7b1c47
feat(compiler): Add support for tiling of element-wise FHELinalg oper…
andidr Apr 9, 2024
35a8703
feat(compiler): Copy tile sizes for all matmul ops when lowering to l…
andidr Apr 9, 2024
865dcfe
feat(compiler): Add support for tiling of fhelinalg.apply_mapped_look…
andidr Apr 9, 2024
a75b5ed
feat(compiler): Add support for tiling of fhelinalg.apply_multi_looku…
andidr Apr 9, 2024
479090e
feat(compiler): Add support for tiling of fhelinalg.sum
andidr Apr 11, 2024
bbd0ff6
refactor(compiler): Remove attribute "tile-sizes" from operations aft…
andidr Apr 11, 2024
a3689ac
feat(compiler): Add support for tiling of fhelinalg.transpose
andidr Apr 11, 2024
1f56008
feat(compiler): Add support for tiling of optimizer.partition_frontier
andidr Apr 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 25 additions & 6 deletions ...lers/concrete-compiler/compiler/lib/Conversion/FHETensorOpsToLinalg/TensorOpsToLinalg.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -305,6 +305,9 @@ struct FHELinalgOpToLinalgGeneric : public mlir::OpRewritePattern<FHELinalgOp> {
ins, outs, maps, iteratorTypes,
doc, call, bodyBuilder);

if (linalgOp->hasAttr("tile-sizes"))
genericOp->setAttr("tile-sizes", linalgOp->getAttr("tile-sizes"));

rewriter.replaceOp(linalgOp, {genericOp.getResult(0)});

return ::mlir::success();
Expand Down Expand Up @@ -438,6 +441,9 @@ struct FHELinalgApplyMappedLookupTableToLinalgGeneric
auto genericOp = rewriter.create<linalg::GenericOp>(
loc, resTys, ins, outs, affineMaps, iteratorTypes, lambdaBlock);

if (mappedLookup->hasAttr("tile-sizes"))
genericOp->setAttr("tile-sizes", mappedLookup->getAttr("tile-sizes"));

rewriter.replaceOp(mappedLookup, {genericOp.getResult(0)});

return ::mlir::success();
Expand Down Expand Up @@ -592,6 +598,9 @@ struct FHELinalgApplyMultiLookupTableToLinalgGeneric
fheLinalgLutOp.getLoc(), resTypes, ins, outs, maps, iteratorTypes,
doc, call, bodyBuilder);

if (fheLinalgLutOp->hasAttr("tile-sizes"))
genericOp->setAttr("tile-sizes", fheLinalgLutOp->getAttr("tile-sizes"));

rewriter.replaceOp(fheLinalgLutOp, {genericOp.getResult(0)});

return ::mlir::success();
Expand Down Expand Up @@ -1153,6 +1162,9 @@ struct SumToLinalgGeneric
linalg::GenericOp genericOp = rewriter.create<linalg::GenericOp>(
location, resultTypes, ins, outs, maps, iteratorTypes, regionBuilder);

if (sumOp->hasAttr("tile-sizes"))
genericOp->setAttr("tile-sizes", sumOp->getAttr("tile-sizes"));

mlir::Value accumulation = genericOp.getResult(0);

mlir::Value result = accumulation;
Expand Down Expand Up @@ -1249,13 +1261,14 @@ struct TransposeToLinalgGeneric
mlir::Value item = blockArgs[0];
nestedBuilder.create<linalg::YieldOp>(location, item);
};
mlir::Value result =
rewriter
.create<linalg::GenericOp>(location, resultTypes, ins, outs, maps,
iteratorTypes, regionBuilder)
.getResult(0);

rewriter.replaceOp(transposeOp, {result});
linalg::GenericOp genericOp = rewriter.create<linalg::GenericOp>(
location, resultTypes, ins, outs, maps, iteratorTypes, regionBuilder);

if (transposeOp->hasAttr("tile-sizes"))
genericOp->setAttr("tile-sizes", transposeOp->getAttr("tile-sizes"));

rewriter.replaceOp(transposeOp, genericOp.getResults());
return mlir::success();
};
};
Expand Down Expand Up @@ -1934,6 +1947,9 @@ struct FHELinalgUnaryOpToLinalgGeneric
ins, outs, maps, iteratorTypes,
doc, call, bodyBuilder);

if (linalgOp->hasAttr("tile-sizes"))
genericOp->setAttr("tile-sizes", linalgOp->getAttr("tile-sizes"));

rewriter.replaceOp(linalgOp, {genericOp.getResult(0)});

return ::mlir::success();
Expand Down Expand Up @@ -2005,6 +2021,9 @@ struct TensorPartitionFrontierOpToLinalgGeneric
outs, maps, iteratorTypes, doc,
call, bodyBuilder);

if (pfOp->hasAttr("tile-sizes"))
genericOp->setAttr("tile-sizes", pfOp->getAttr("tile-sizes"));

rewriter.replaceOp(pfOp, {genericOp.getResult(0)});

return ::mlir::success();
Expand Down
24 changes: 22 additions & 2 deletions compilers/concrete-compiler/compiler/lib/Dialect/FHELinalg/Transforms/Tiling.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,7 @@ class GenericTilingPattern
if (lres.succeeded()) {
res.value().tileOp->setAttr(kTransformMarker, rewriter.getUnitAttr());
res.value().tiledOp->setAttr(kTransformMarker, rewriter.getUnitAttr());
res.value().tiledOp->removeAttr("tile-sizes");
rewriter.replaceOp(op.getOperation(), res.value().tileOp->getResults());
}

Expand Down Expand Up @@ -141,6 +142,7 @@ class GenericTilingPattern
if (lres.succeeded()) {
res.value().parallelTiledOp->setAttr(kTransformMarker,
rewriter.getUnitAttr());
res.value().parallelTiledOp->removeAttr("tile-sizes");
res.value().mergeOp->setAttr(kTransformMarker, rewriter.getUnitAttr());
res.value().initialOp->setAttr(kTransformMarker,
rewriter.getUnitAttr());
Expand Down Expand Up @@ -186,9 +188,27 @@ class FHELinalgTilingMarkerPass
mlir::Builder(&this->getContext()).getI64ArrayAttr(tileSizes);

op->walk([&](mlir::Operation *op) {
if (llvm::isa<mlir::concretelang::FHELinalg::ApplyLookupTableEintOp>(
if (llvm::isa<mlir::concretelang::FHELinalg::ApplyLookupTableEintOp,
mlir::concretelang::FHELinalg::ApplyMappedLookupTableEintOp,
mlir::concretelang::FHELinalg::ApplyMultiLookupTableEintOp>(
op) ||
llvm::isa<mlir::concretelang::FHELinalg::MatMulEintIntOp>(op)) {
llvm::isa<mlir::concretelang::FHELinalg::MatMulEintIntOp,
mlir::concretelang::FHELinalg::MatMulIntEintOp,
mlir::concretelang::FHELinalg::MatMulEintEintOp>(op) ||
llvm::isa<mlir::concretelang::FHELinalg::SumOp>(op) ||
llvm::isa<mlir::concretelang::FHELinalg::TransposeOp>(op) ||
llvm::isa<mlir::concretelang::FHELinalg::AddEintOp,
mlir::concretelang::FHELinalg::AddEintIntOp,
mlir::concretelang::FHELinalg::SubIntEintOp,
mlir::concretelang::FHELinalg::SubEintIntOp,
mlir::concretelang::FHELinalg::SubEintOp,
mlir::concretelang::FHELinalg::MulEintIntOp,
mlir::concretelang::FHELinalg::MulEintOp,
mlir::concretelang::FHELinalg::LsbEintOp,
mlir::concretelang::FHELinalg::NegEintOp,
mlir::concretelang::FHELinalg::ToSignedOp,
mlir::concretelang::FHELinalg::ToUnsignedOp,
mlir::concretelang::FHELinalg::RoundOp>(op)) {
op->setAttr("tile-sizes", tileAttr);
}
});
Expand Down
93 changes: 90 additions & 3 deletions compilers/concrete-compiler/compiler/tests/check_tests/Dialect/FHELinalg/tiling.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
// CHECK-NEXT: %[[V7:.*]] = affine.apply #map1(%[[Varg2]])
// CHECK-NEXT: %[[Vextracted_slice_0:.*]] = tensor.extract_slice %[[Varg0:.*]]{{\[0,}} %[[V6]]{{\] \[8, 2\] \[1, 1\]}} : tensor<8x4x!FHE.eint<6>> to tensor<8x2x!FHE.eint<6>>
// CHECK-NEXT: %[[Vextracted_slice_1:.*]] = tensor.extract_slice %[[Varg1:.*]]{{\[}}%[[V7]], 0{{\] \[2, 2\] \[1, 1\]}} : tensor<4x2xi7> to tensor<2x2xi7>
// CHECK-NEXT: %[[V8:.*]] = linalg.generic {indexing_maps = {{\[}}#map2, #map3, #map4{{\], iterator}}_types = {{\[}}"parallel", "parallel", "reduction"{{\]}}} ins(%[[Vextracted_slice_0]], %[[Vextracted_slice_1]] : tensor<8x2x!FHE.eint<6>>, tensor<2x2xi7>) outs(%[[Vextracted_slice]] : tensor<8x2x!FHE.eint<6>>) attrs = {"tile-sizes" = {{\[0, 0, 2\]}}} {
// CHECK-NEXT: %[[V8:.*]] = linalg.generic {indexing_maps = {{\[}}#map2, #map3, #map4{{\], iterator}}_types = {{\[}}"parallel", "parallel", "reduction"{{\]}}} ins(%[[Vextracted_slice_0]], %[[Vextracted_slice_1]] : tensor<8x2x!FHE.eint<6>>, tensor<2x2xi7>) outs(%[[Vextracted_slice]] : tensor<8x2x!FHE.eint<6>>) {
// CHECK-NEXT: ^bb0(%[[Vin:.*]]: !FHE.eint<6>, %[[Vin_2:.*]]: i7, %[[Vout:.*]]: !FHE.eint<6>):
// CHECK-NEXT: %[[V9:.*]] = "FHE.mul_eint_int"(%[[Vin]], %[[Vin_2]]) : (!FHE.eint<6>, i7) -> !FHE.eint<6>
// CHECK-NEXT: %[[V10:.*]] = "FHE.add_eint"(%[[Vout]], %[[V9]]) : (!FHE.eint<6>, !FHE.eint<6>) -> !FHE.eint<6>
Expand All @@ -30,7 +30,7 @@ func.func @tiled_2(%a: tensor<8x4x!FHE.eint<6>>, %b: tensor<4x2xi7>) -> tensor<8
// CHECK-NEXT: %[[V7:.*]] = affine.apply #map1(%[[Varg2]])
// CHECK-NEXT: %[[Vextracted_slice_0:.*]] = tensor.extract_slice %[[Varg0:.*]]{{\[0,}} %[[V6]]{{\] \[8, 4\] \[1, 1\]}} : tensor<8x4x!FHE.eint<6>> to tensor<8x4x!FHE.eint<6>>
// CHECK-NEXT: %[[Vextracted_slice_1:.*]] = tensor.extract_slice %[[Varg1:.*]]{{\[}}%[[V7]], 0{{\] \[4, 2\] \[1, 1\]}} : tensor<4x2xi7> to tensor<4x2xi7>
// CHECK-NEXT: %[[V8:.*]] = linalg.generic {indexing_maps = {{\[}}#map2, #map3, #map4{{\], iterator}}_types = {{\[}}"parallel", "parallel", "reduction"{{\]}}} ins(%[[Vextracted_slice_0]], %[[Vextracted_slice_1]] : tensor<8x4x!FHE.eint<6>>, tensor<4x2xi7>) outs(%[[Vextracted_slice]] : tensor<8x2x!FHE.eint<6>>) attrs = {"tile-sizes" = {{\[0, 0, 4\]}}} {
// CHECK-NEXT: %[[V8:.*]] = linalg.generic {indexing_maps = {{\[}}#map2, #map3, #map4{{\], iterator}}_types = {{\[}}"parallel", "parallel", "reduction"{{\]}}} ins(%[[Vextracted_slice_0]], %[[Vextracted_slice_1]] : tensor<8x4x!FHE.eint<6>>, tensor<4x2xi7>) outs(%[[Vextracted_slice]] : tensor<8x2x!FHE.eint<6>>) {
// CHECK-NEXT: ^bb0(%[[Vin:.*]]: !FHE.eint<6>, %[[Vin_2:.*]]: i7, %[[Vout:.*]]: !FHE.eint<6>):
// CHECK-NEXT: %[[V9:.*]] = "FHE.mul_eint_int"(%[[Vin]], %[[Vin_2]]) : (!FHE.eint<6>, i7) -> !FHE.eint<6>
// CHECK-NEXT: %[[V10:.*]] = "FHE.add_eint"(%[[Vout]], %[[V9]]) : (!FHE.eint<6>, !FHE.eint<6>) -> !FHE.eint<6>
Expand All @@ -57,7 +57,7 @@ func.func @tiled_one_big_tile(%a: tensor<8x4x!FHE.eint<6>>, %b: tensor<4x2xi7>)
// CHECK-NEXT: %[[V7:.*]] = affine.apply #map(%[[Varg4]])
// CHECK-NEXT: %[[Vextracted_slice:.*]] = tensor.extract_slice %[[Varg0]]{{\[}}%[[V2]], %[[V3]], %[[V4]]{{\] \[2, 3, 2\] \[1, 1, 1\]}} : tensor<2x3x4x!FHE.eint<2>> to tensor<2x3x2x!FHE.eint<2>>
// CHECK-NEXT: %[[Vextracted_slice_0:.*]] = tensor.extract_slice %[[Varg5]]{{\[}}%[[V5]], %[[V6]], %[[V7]]{{\] \[2, 3, 2\] \[1, 1, 1\]}} : tensor<2x3x4x!FHE.eint<2>> to tensor<2x3x2x!FHE.eint<2>>
// CHECK-NEXT: %[[V8:.*]] = linalg.generic {indexing_maps = {{\[}}#map2, #map2{{\], iterator}}_types = {{\[}}"parallel", "parallel", "parallel"{{\]}}} ins(%[[Vextracted_slice]] : tensor<2x3x2x!FHE.eint<2>>) outs(%[[Vextracted_slice_0]] : tensor<2x3x2x!FHE.eint<2>>) attrs = {"tile-sizes" = {{\[2, 3, 2\]}}} {
// CHECK-NEXT: %[[V8:.*]] = linalg.generic {indexing_maps = {{\[}}#map2, #map2{{\], iterator}}_types = {{\[}}"parallel", "parallel", "parallel"{{\]}}} ins(%[[Vextracted_slice]] : tensor<2x3x2x!FHE.eint<2>>) outs(%[[Vextracted_slice_0]] : tensor<2x3x2x!FHE.eint<2>>) {
// CHECK-NEXT: ^bb0(%[[Vin:.*]]: !FHE.eint<2>, %[[Vout:.*]]: !FHE.eint<2>):
// CHECK-NEXT: %[[V12:.*]] = "FHE.apply_lookup_table"(%[[Vin]], %[[Varg1]]) : (!FHE.eint<2>, tensor<4xi64>) -> !FHE.eint<2>
// CHECK-NEXT: linalg.yield %[[V12]] : !FHE.eint<2>
Expand All @@ -73,3 +73,90 @@ func.func @apply_lookup_table(%arg0: tensor<2x3x4x!FHE.eint<2>>, %arg1: tensor<4
%1 = "FHELinalg.apply_lookup_table"(%arg0, %arg1) { "tile-sizes" = [2,3,2] } : (tensor<2x3x4x!FHE.eint<2>>, tensor<4xi64>) -> (tensor<2x3x4x!FHE.eint<2>>)
return %1: tensor<2x3x4x!FHE.eint<2>>
}

// -----

// CHECK: %[[res:.*]] = linalg.generic {indexing_maps = [#[[map:.*]], #[[map2:.*]], #[[map3:.*]]], iterator_types = ["parallel", "parallel", "parallel"]} ins(%[[extracted_slice:.*]], %[[extracted_slice_0:.*]] : tensor<2x3x2x!FHE.eint<7>>, tensor<2x3x2xindex>) outs(%[[extracted_slice_1:.*]] : tensor<2x3x2x!FHE.eint<7>>)
func.func @apply_mapped_lookup_table(
%input: tensor<2x3x4x!FHE.eint<7>>,
%luts: tensor<10x128xi64>,
%map: tensor<2x3x4xindex>
) -> tensor<2x3x4x!FHE.eint<7>> {
%0 = "FHELinalg.apply_mapped_lookup_table"(%input, %luts, %map) { "tile-sizes" = [2,3,2] } : (tensor<2x3x4x!FHE.eint<7>>, tensor<10x128xi64>, tensor<2x3x4xindex>) -> (tensor<2x3x4x!FHE.eint<7>>)
return %0: tensor<2x3x4x!FHE.eint<7>>
}

// -----

// CHECK: %[[res:.*]] = linalg.generic {indexing_maps = [#[[map:.*]], #[[map2:.*]]], iterator_types = ["parallel", "parallel"]} ins(%[[extracted_slice:.*]] : tensor<3x3x!FHE.eint<2>>) outs(%[[extracted_slice_0:.*]] : tensor<3x3x!FHE.eint<2>>) {
func.func @main(%arg0: tensor<3x3x!FHE.eint<2>>, %arg1: tensor<3x3x4xi8>) -> tensor<3x3x!FHE.eint<2>> {
%1 = "FHELinalg.apply_multi_lookup_table"(%arg0, %arg1) { "tile-sizes" = [3, 3] }: (tensor<3x3x!FHE.eint<2>>, tensor<3x3x4xi8>) -> tensor<3x3x!FHE.eint<2>>
return %1: tensor<3x3x!FHE.eint<2>>
}

// -----

// CHECK: %[[V4:.*]] = scf.forall (%[[Varg1]], %[[Varg2]], %[[Varg3]]) in (1, 1, 2) shared_outs(%[[Varg4:.*]] = %[[V3]]) -> (tensor<3x8x2x!FHE.eint<7>>) {
// CHECK-NEXT: %[[Vextracted_slice:.*]] = tensor.extract_slice %[[Varg4]]{{\[0, 0,}} %[[Varg1]]{{\] \[3, 8, 1\] \[1, 1, 1\]}} : tensor<3x8x2x!FHE.eint<7>> to tensor<3x8x!FHE.eint<7>>
// CHECK-NEXT: %[[V6:.*]] = affine.apply #map1(%[[Varg1]])
// CHECK-NEXT: %[[V7:.*]] = affine.apply #map2(%[[Varg2]])
// CHECK-NEXT: %[[V8:.*]] = affine.apply #map3(%[[Varg3]])
// CHECK-NEXT: %[[V9:.*]] = affine.apply #map1(%[[Varg1]])
// CHECK-NEXT: %[[V10:.*]] = affine.apply #map2(%[[Varg2]])
// CHECK-NEXT: %[[Vextracted_slice_0:.*]] = tensor.extract_slice %[[Varg0]]{{\[}}%[[V6]], %[[V7]], %[[V8]]{{\] \[3, 8, 2\] \[1, 1, 1\]}} : tensor<3x8x4x!FHE.eint<7>> to tensor<3x8x2x!FHE.eint<7>>
// CHECK-NEXT: %[[Vextracted_slice_1:.*]] = tensor.extract_slice %[[Vextracted_slice]]{{\[}}%[[V9]], %[[V10]]{{\] \[3, 8\] \[1, 1\]}} : tensor<3x8x!FHE.eint<7>> to tensor<3x8x!FHE.eint<7>>
// CHECK-NEXT: %[[V11:.*]] = linalg.generic {indexing_maps = {{\[}}#map, #map4{{\], iterator}}_types = {{\[}}"parallel", "parallel", "reduction"{{\]}}} ins(%[[Vextracted_slice_0]] : tensor<3x8x2x!FHE.eint<7>>) outs(%[[Vextracted_slice_1]] : tensor<3x8x!FHE.eint<7>>) {
// CHECK-NEXT: ^bb0(%[[Vin:.*]]: !FHE.eint<7>, %[[Vout:.*]]: !FHE.eint<7>):
// CHECK-NEXT: %[[V14:.*]] = "FHE.add_eint"(%[[Vin]], %[[Vout]]) : (!FHE.eint<7>, !FHE.eint<7>) -> !FHE.eint<7>
// CHECK-NEXT: linalg.yield %[[V14]] : !FHE.eint<7>
// CHECK-NEXT: } -> tensor<3x8x!FHE.eint<7>>
// CHECK-NEXT: %[[V12:.*]] = affine.apply #map1(%[[Varg1]])
// CHECK-NEXT: %[[V13:.*]] = affine.apply #map2(%[[Varg2]])
// CHECK-NEXT: scf.forall.in_parallel {
// CHECK-NEXT: tensor.parallel_insert_slice %[[V11]] into %[[Varg4]]{{\[}}%[[V12]], %[[V13]], %[[Varg1]]{{\] \[3, 8, 1\] \[1, 1, 1\]}} : tensor<3x8x!FHE.eint<7>> into tensor<3x8x2x!FHE.eint<7>>
// CHECK-NEXT: }
// CHECK-NEXT: }

func.func @main(%arg0: tensor<3x8x4x!FHE.eint<7>>) -> tensor<3x8x!FHE.eint<7>> {
%0 = "FHELinalg.sum"(%arg0) { axes = [2], "tile-sizes" = [3, 8, 2] } : (tensor<3x8x4x!FHE.eint<7>>) -> tensor<3x8x!FHE.eint<7>>
return %0 : tensor<3x8x!FHE.eint<7>>
}

// -----

// CHECK: #map = affine_map<(d0) -> (d0 * -3 + 2, 3)>
// CHECK-NEXT: #map1 = affine_map<(d0) -> (d0 * -8 + 10, 8)>
// CHECK-NEXT: #map2 = affine_map<(d0) -> (d0 * 3)>
// CHECK-NEXT: #map3 = affine_map<(d0) -> (d0 * 8)>
// CHECK-NEXT: #map4 = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-NEXT: #map5 = affine_map<(d0, d1) -> (d1, d0)>
// CHECK-NEXT: module {
// CHECK-NEXT: func.func @transpose_eint_2D(%[[Varg0:.*]]: tensor<2x10x!FHE.eint<6>>) -> tensor<10x2x!FHE.eint<6>> {
// CHECK-NEXT: %[[V0:.*]] = "FHE.zero_tensor"() : () -> tensor<10x2x!FHE.eint<6>>
// CHECK-NEXT: %[[V1:.*]] = scf.forall (%[[Varg1]], %[[Varg2]]) in (1, 2) shared_outs(%[[Varg3:.*]] = %[[V0]]) -> (tensor<10x2x!FHE.eint<6>>) {
// CHECK-NEXT: %[[V2:.*]] = affine.min #map(%[[Varg1]])
// CHECK-NEXT: %[[V3:.*]] = affine.min #map1(%[[Varg2]])
// CHECK-NEXT: %[[V4:.*]] = affine.apply #map2(%[[Varg1]])
// CHECK-NEXT: %[[V5:.*]] = affine.apply #map3(%[[Varg2]])
// CHECK-NEXT: %[[V6:.*]] = affine.apply #map3(%[[Varg2]])
// CHECK-NEXT: %[[V7:.*]] = affine.apply #map2(%[[Varg1]])
// CHECK-NEXT: %[[Vextracted_slice:.*]] = tensor.extract_slice %[[Varg0]]{{\[}}%[[V4]], %[[V5]]{{\] \[}}%[[V2]], %[[V3]]{{\] \[1, 1\]}} : tensor<2x10x!FHE.eint<6>> to tensor<?x?x!FHE.eint<6>>
// CHECK-NEXT: %[[Vextracted_slice_0:.*]] = tensor.extract_slice %[[Varg3]]{{\[}}%[[V6]], %[[V7]]{{\] \[}}%[[V3]], %[[V2]]{{\] \[1, 1\]}} : tensor<10x2x!FHE.eint<6>> to tensor<?x?x!FHE.eint<6>>
// CHECK-NEXT: %[[V8:.*]] = linalg.generic {indexing_maps = {{\[}}#map4, #map5{{\], iterator}}_types = {{\[}}"parallel", "parallel"{{\]}}} ins(%[[Vextracted_slice]] : tensor<?x?x!FHE.eint<6>>) outs(%[[Vextracted_slice_0]] : tensor<?x?x!FHE.eint<6>>) {
// CHECK-NEXT: ^bb0(%[[Vin:.*]]: !FHE.eint<6>, %[[Vout:.*]]: !FHE.eint<6>):
// CHECK-NEXT: linalg.yield %[[Vin]] : !FHE.eint<6>
// CHECK-NEXT: } -> tensor<?x?x!FHE.eint<6>>
// CHECK-NEXT: %[[V9:.*]] = affine.apply #map3(%[[Varg2]])
// CHECK-NEXT: %[[V10:.*]] = affine.apply #map2(%[[Varg1]])
// CHECK-NEXT: scf.forall.in_parallel {
// CHECK-NEXT: tensor.parallel_insert_slice %[[V8]] into %[[Varg3]]{{\[}}%[[V9]], %[[V10]]{{\] \[}}%[[V3]], %[[V2]]{{\] \[1, 1\]}} : tensor<?x?x!FHE.eint<6>> into tensor<10x2x!FHE.eint<6>>
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: return %[[V1]] : tensor<10x2x!FHE.eint<6>>
// CHECK-NEXT: }
// CHECK-NEXT: }

func.func @transpose_eint_2D(%arg0: tensor<2x10x!FHE.eint<6>>) -> tensor<10x2x!FHE.eint<6>> {
%c = "FHELinalg.transpose"(%arg0) { "tile-sizes" = [3, 8, 2] } : (tensor<2x10x!FHE.eint<6>>) -> tensor<10x2x!FHE.eint<6>>
return %c : tensor<10x2x!FHE.eint<6>>
}
Loading