Improve single element tensor handling

lib/Transforms/ConvertToCiphertextSemantics/AssignLayout.cpp

@@ -68,8 +68,27 @@ static FailureOr<Value> implementUnpackOpStep(
 
   RankedTensorType unpackedTensorType = dyn_cast<RankedTensorType>(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<int64_t> point = anyRangePoint(rel);
     if (point.empty()) {
       return builder.emitError()
@@ -83,15 +102,16 @@ static FailureOr<Value> 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<int> domainSchedule;
@@ -346,11 +366,15 @@ static FailureOr<Value> 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<Value> implementAssignLayoutStep(
       createdOpCallback(splatOp);
       return splatOp.getResult();
     }
-    auto constantOp = arith::ConstantOp::create(
-        builder, targetType,
-        SplatElementsAttr::get(targetType,
-                               splatAttr.getSplatValue<TypedAttr>()));
-    createdOpCallback(constantOp);
-    return constantOp.getResult();
+    if (inputIsSplatConstant) {
+      auto constantOp = arith::ConstantOp::create(
+          builder, targetType,
+          SplatElementsAttr::get(targetType,
+                                 splatAttr.getSplatValue<TypedAttr>()));
+      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<Value> 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;

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<1xi16>> {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<1xi16>> {tensor_ext.layout = #dense_layout})
+    // CHECK: return
+    return %0 : !secret.secret<tensor<1xi16>>
+  }
+}
+
+// -----
+
+// 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<1xi16>> {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<1xi16>> {tensor_ext.layout = #sparse_layout})
+    // CHECK: return
+    return %0 : !secret.secret<tensor<1xi16>>
+  }
+}

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<originalType = tensor<1xi16>, 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<originalType = tensor<1xi16>, 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>
+}