From 7c361bcb2978d21effce8120948c562dbe0ea4ce Mon Sep 17 00:00:00 2001 From: Marc Desgroseilliers Date: Wed, 5 Aug 2026 09:32:50 +0200 Subject: [PATCH] handle single element tensors fold single element tensor in existing if --- .../AssignLayout.cpp | 75 ++++++++++++++----- .../single_element_assign_layout.mlir | 45 +++++++++++ .../Transforms/lower_unpack/lower_unpack.mlir | 33 ++++++++ 3 files changed, 135 insertions(+), 18 deletions(-) create mode 100644 tests/Transforms/convert_to_ciphertext_semantics/single_element_assign_layout.mlir diff --git a/lib/Transforms/ConvertToCiphertextSemantics/AssignLayout.cpp b/lib/Transforms/ConvertToCiphertextSemantics/AssignLayout.cpp index a2ee919552..c2419722d7 100644 --- a/lib/Transforms/ConvertToCiphertextSemantics/AssignLayout.cpp +++ b/lib/Transforms/ConvertToCiphertextSemantics/AssignLayout.cpp @@ -68,8 +68,27 @@ static FailureOr implementUnpackOpStep( RankedTensorType unpackedTensorType = dyn_cast(targetType); - if (!unpackedTensorType) { - // it's a scalar, so we can extract from any slot in the mapping + // Restrict the layout relation's domain bounds to the valid elements + // [0, dimSize - 1] of targetType. This has to happen before any use of `rel` + // below, including the single-element shortcut: that shortcut samples an + // arbitrary point from the relation's range, and with an unrestricted domain + // the sample can be the slot belonging to a different element. + if (unpackedTensorType) { + for (unsigned i = 0; i < unpackedTensorType.getRank(); ++i) { + if (unpackedTensorType.isDynamicDim(i)) continue; + rel.addBound(presburger::BoundType::UB, + rel.getVarKindOffset(presburger::VarKind::Domain) + i, + unpackedTensorType.getDimSize(i) - 1); + } + } + + bool isSingleElementTensor = unpackedTensorType && + unpackedTensorType.hasStaticShape() && + unpackedTensorType.getNumElements() == 1; + bool isScalar = !unpackedTensorType; + + if (isScalar || isSingleElementTensor) { + // Extract the lone element from any slot in the mapping. std::vector point = anyRangePoint(rel); if (point.empty()) { return builder.emitError() @@ -83,15 +102,16 @@ static FailureOr implementUnpackOpStep( } auto extractOp = tensor::ExtractOp::create(builder, input, indices); createdOpCallback(extractOp); - return extractOp.getResult(); - } - // Restrict the layout relation's domain bounds to the valid elements - // [0, dimSize - 1] of targetType. - for (unsigned i = 0; i < unpackedTensorType.getRank(); ++i) { - if (unpackedTensorType.isDynamicDim(i)) continue; - rel.addBound(presburger::BoundType::UB, i, - unpackedTensorType.getDimSize(i) - 1); + if (isScalar) { + return extractOp.getResult(); + } + + // isSingleElementTensor -> wrap back into single-element tensor. + auto fromElementsOp = tensor::FromElementsOp::create( + builder, unpackedTensorType, ValueRange{extractOp.getResult()}); + createdOpCallback(fromElementsOp); + return fromElementsOp.getResult(); } SmallVector domainSchedule; @@ -346,11 +366,15 @@ static FailureOr implementAssignLayoutStep( } // The result can be simplified if the layout is dense in the ciphertext type, - // and the input is a scalar or a constant splat. + // and the input is a scalar, a constant splat, or a single-element tensor + // (whose dense packing is a broadcast of its lone element). SplatElementsAttr splatAttr; bool inputIsScalar = !dataSemanticType; bool inputIsSplatConstant = matchPattern(input, m_Constant(&splatAttr)); - if ((inputIsScalar || inputIsSplatConstant) && + bool inputIsSingleElementTensor = dataSemanticType && + dataSemanticType.hasStaticShape() && + dataSemanticType.getNumElements() == 1; + if ((inputIsScalar || inputIsSplatConstant || inputIsSingleElementTensor) && isDenseLayout(rel, targetType)) { // Regardless of being constant or not, a scalar can be splat into the // ciphertext tensor. @@ -359,12 +383,27 @@ static FailureOr implementAssignLayoutStep( createdOpCallback(splatOp); return splatOp.getResult(); } - auto constantOp = arith::ConstantOp::create( - builder, targetType, - SplatElementsAttr::get(targetType, - splatAttr.getSplatValue())); - createdOpCallback(constantOp); - return constantOp.getResult(); + if (inputIsSplatConstant) { + auto constantOp = arith::ConstantOp::create( + builder, targetType, + SplatElementsAttr::get(targetType, + splatAttr.getSplatValue())); + createdOpCallback(constantOp); + return constantOp.getResult(); + } + // A non-constant single-element tensor: the loop-generator path below + // handles it correctly, but emits one iteration per slot whose body is + // loop-invariant; extract the element once and splat it instead. + auto zero = arith::ConstantIndexOp::create(builder, 0); + createdOpCallback(zero); + SmallVector zeroIndices(dataSemanticType.getRank(), + zero.getResult()); + auto extractOp = tensor::ExtractOp::create(builder, input, zeroIndices); + createdOpCallback(extractOp); + auto splatOp = + tensor::SplatOp::create(builder, targetType, extractOp.getResult()); + createdOpCallback(splatOp); + return splatOp.getResult(); } DenseElementsAttr constantAttr; diff --git a/tests/Transforms/convert_to_ciphertext_semantics/single_element_assign_layout.mlir b/tests/Transforms/convert_to_ciphertext_semantics/single_element_assign_layout.mlir new file mode 100644 index 0000000000..40d6c37a60 --- /dev/null +++ b/tests/Transforms/convert_to_ciphertext_semantics/single_element_assign_layout.mlir @@ -0,0 +1,45 @@ +// RUN: heir-opt %s --convert-to-ciphertext-semantics=min-slot-count=32 --split-input-file | FileCheck %s + +// Assigning a layout to a single-element tensor whose layout is dense in the +// ciphertext is a broadcast of its lone element: extract the element once and +// splat it, rather than emitting one loop iteration per slot whose body is +// loop-invariant. +#dense_layout = #tensor_ext.layout<"{ [i0] -> [ct, slot] : ct = 0 and i0 = 0 and 0 <= slot <= 31 }"> + +// CHECK: @assign_layout_single_element_dense +module { + func.func @assign_layout_single_element_dense(%arg0: tensor<1xi16>) -> (!secret.secret> {tensor_ext.layout = #dense_layout}) { + // CHECK-NOT: scf.for + // CHECK-NOT: tensor.insert + // CHECK: %[[ELT:.*]] = tensor.extract + // CHECK: tensor.splat %[[ELT]] : tensor<1x32xi16> + // CHECK-NOT: scf.for + // CHECK-NOT: tensor.insert + %0 = secret.generic() { + %1 = tensor_ext.assign_layout %arg0 {layout = #dense_layout, tensor_ext.layout = #dense_layout} : tensor<1xi16> + secret.yield %1 : tensor<1xi16> + } -> (!secret.secret> {tensor_ext.layout = #dense_layout}) + // CHECK: return + return %0 : !secret.secret> + } +} + +// ----- + +// When the layout is not dense in the ciphertext, the lone element does not +// occupy every slot, so the broadcast above would be wrong. This falls back to +// the general loop-generator path. +#sparse_layout = #tensor_ext.layout<"{ [i0] -> [ct, slot] : ct = 0 and i0 = 0 and 0 <= slot <= 15 }"> + +// CHECK: @assign_layout_single_element_not_dense +module { + func.func @assign_layout_single_element_not_dense(%arg0: tensor<1xi16>) -> (!secret.secret> {tensor_ext.layout = #sparse_layout}) { + // CHECK-NOT: tensor.splat + %0 = secret.generic() { + %1 = tensor_ext.assign_layout %arg0 {layout = #sparse_layout, tensor_ext.layout = #sparse_layout} : tensor<1xi16> + secret.yield %1 : tensor<1xi16> + } -> (!secret.secret> {tensor_ext.layout = #sparse_layout}) + // CHECK: return + return %0 : !secret.secret> + } +} diff --git a/tests/Transforms/lower_unpack/lower_unpack.mlir b/tests/Transforms/lower_unpack/lower_unpack.mlir index 3d104c2cdc..ba7c09397d 100644 --- a/tests/Transforms/lower_unpack/lower_unpack.mlir +++ b/tests/Transforms/lower_unpack/lower_unpack.mlir @@ -38,3 +38,36 @@ func.func @unpack_rotated_tensor2(%arg0: tensor<64x64xi16> {tensor_ext.original_ %0 = tensor_ext.unpack %arg0 {layout=#tensor_layout2} : (tensor<64x64xi16>) -> tensor<32xi16> return %0 : tensor<32xi16> } + +// A single-element result is constant work: read the lone element once rather +// than looping over the ciphertext's slots and accumulating into the result. +#singleton_layout = #tensor_ext.layout<"{ [i0] -> [ct, slot] : ct = 0 and i0 = 0 and 0 <= slot <= 1023 }"> +#singleton_original_type = #tensor_ext.original_type, layout = #singleton_layout> + +// CHECK: @unpack_single_element +// CHECK-NOT: scf.for +// CHECK-NOT: tensor.insert +// CHECK: tensor.extract +// CHECK: tensor.from_elements +// CHECK: return +func.func @unpack_single_element(%arg0: tensor<1x1024xi16> {tensor_ext.original_type = #singleton_original_type}) -> tensor<1xi16> { + %0 = tensor_ext.unpack %arg0 {layout=#singleton_layout} : (tensor<1x1024xi16>) -> tensor<1xi16> + return %0 : tensor<1xi16> +} + +// The relation's domain (0 <= i0 <= 31) is wider than the single-element result +// type, and each element lives in exactly one slot. The mapping is rotated by 5 +// so element 0 sits at slot 5, away from either end of the slot range: the slot +// to read must be picked *after* the domain is clamped to the result's extent, +// because sampling the unclamped relation returns the slot of some other +// element -- silently the wrong value rather than invalid IR. +#loose_layout = #tensor_ext.layout<"{ [i0] -> [ct, slot] : ct = 0 and (slot - i0 - 5) mod 32 = 0 and 0 <= i0 <= 31 and 0 <= slot <= 31 }"> +#loose_original_type = #tensor_ext.original_type, layout = #loose_layout> + +// CHECK: @unpack_single_element_loose_domain +// CHECK-DAG: %[[C5:.*]] = arith.constant 5 : index +// CHECK: tensor.extract %{{[^[]*}}[%{{[^,]*}}, %[[C5]]] +func.func @unpack_single_element_loose_domain(%arg0: tensor<1x1024xi16> {tensor_ext.original_type = #loose_original_type}) -> tensor<1xi16> { + %0 = tensor_ext.unpack %arg0 {layout=#loose_layout} : (tensor<1x1024xi16>) -> tensor<1xi16> + return %0 : tensor<1xi16> +}