Merge pull request #2653 from VedantParanjape:permutation

PiperOrigin-RevId: 871950476

Author: copybara-github

lib/Dialect/TensorExt/IR/TensorExtOps.cpp

@@ -129,6 +129,24 @@ LogicalResult verifyLayoutMatchesType(const Attribute& layoutAttr, Type type,
     return success();
   }
 
+  if (auto denseElementsAttr = dyn_cast<DenseIntElementsAttr>(layoutAttr)) {
+    // Assert the attr has shape <N x 4>
+    int64_t rank = denseElementsAttr.getType().getRank();
+    if (rank != 2)
+      return op->emitOpError()
+             << "requires permutation attribute to Rank 2, but "
+             << "found shape <" << denseElementsAttr.getType() << ">";
+
+    int64_t cols = denseElementsAttr.getType().getDimSize(1);
+    if (cols != 4)
+      return op->emitOpError()
+             << "requires permutation attribute to be of shape <N x 4>, but "
+                "found shape <"
+             << denseElementsAttr.getType() << ">" << "Rank: " << rank
+             << " Cols: " << cols << "\n";
+    return success();
+  }
+
   return op->emitOpError("Unsupported layout attribute");
 }
 

lib/Dialect/TensorExt/IR/TensorExtOps.td

@@ -92,6 +92,9 @@ def TensorExt_RemapOp : TensorExt_Op<"remap", [Pure, AllTypesMatch<["input", "ou
 // Forces ops to use a general Attribute and dyn_cast to the specific kind of
 // layout they support.
 def LayoutLike : AnyAttrOf<[
+  // A list of tuples [a, b, c, d] representing an explicit map (ct, slot) ->
+  // (ct, slot) defined by f(a, b) = (c, d).
+  AnyI64ElementsAttr,
   TensorExt_LayoutAttr,
 ]>;
 

lib/Transforms/ConvertToCiphertextSemantics/AssignLayout.cpp

@@ -195,6 +195,51 @@ static FailureOr<Value> implementUnpackOpNew(
   return loop.value().getResults()[0];
 }
 
+static FailureOr<Value> implementAssignLayoutPermutation(
+    Value input, DenseIntElementsAttr permutation, int64_t ciphertextSize,
+    ImplicitLocOpBuilder& builder,
+    const std::function<void(Operation*)>& createdOpCallback) {
+  auto elementType = getElementTypeOrSelf(input.getType());
+
+  // TODO(#2666): This logic assumes ctxt = 0, this can be extended to handle a
+  // more general case. permutation is <N x 4 x i64>: each row is [src_ct,
+  // src_slot, dst_ct, dst_slot]. src_ct and dst_ct are always 0 (single
+  // ciphertext), so treat as 1D: extract from input[0][src_slot] and insert
+  // into result[0][dst_slot].
+  auto ciphertextType = RankedTensorType::get({1, ciphertextSize}, elementType);
+  auto zeroCtxt = arith::ConstantOp::create(
+      builder, ciphertextType, builder.getZeroAttr(ciphertextType));
+  createdOpCallback(zeroCtxt);
+  Value result = zeroCtxt.getResult();
+  auto ctIdxOp = arith::ConstantIndexOp::create(builder, 0);
+  createdOpCallback(ctIdxOp);
+  Value ctIdx = ctIdxOp.getResult();
+
+  for (auto it = permutation.value_begin<APInt>();
+       it != permutation.value_end<APInt>();) {
+    // skip src_ct (always 0)
+    ++it;
+    int64_t srcSlot = (*it++).getSExtValue();
+    // skip dst_ct (always 0)
+    ++it;
+    int64_t dstSlot = (*it++).getSExtValue();
+
+    auto srcSlotIdx = arith::ConstantIndexOp::create(builder, srcSlot);
+    createdOpCallback(srcSlotIdx);
+    auto extracted = tensor::ExtractOp::create(builder, input,
+                                               ValueRange{ctIdx, srcSlotIdx});
+    createdOpCallback(extracted);
+    auto dstSlotIdx = arith::ConstantIndexOp::create(builder, dstSlot);
+    createdOpCallback(dstSlotIdx);
+    auto insertOp = tensor::InsertOp::create(
+        builder, extracted.getResult(), result, ValueRange{ctIdx, dstSlotIdx});
+    createdOpCallback(insertOp);
+    result = insertOp.getResult();
+  }
+
+  return result;
+}
+
 FailureOr<Value> implementAssignLayout(
     Value input, Attribute layout, int64_t ciphertextSize,
     ImplicitLocOpBuilder& builder,
@@ -203,6 +248,10 @@ FailureOr<Value> implementAssignLayout(
   if (LayoutAttr layoutAttr = dyn_cast<LayoutAttr>(layout)) {
     return implementAssignLayoutNew(input, layoutAttr, ciphertextSize, builder,
                                     createdOpCallback);
+  } else if (DenseIntElementsAttr elementAttr =
+                 dyn_cast<DenseIntElementsAttr>(layout)) {
+    return implementAssignLayoutPermutation(input, elementAttr, ciphertextSize,
+                                            builder, createdOpCallback);
   }
   return builder.emitError() << "Unsupported layout attribute type: " << layout;
 };

lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp

@@ -206,6 +206,17 @@ struct LayoutMaterializationTypeConverter
         [this](IndexType type, LayoutAttr attr) -> std::optional<Type> {
           return materializeLayout(type, attr, getCiphertextSize());
         });
+    addConversion([this](secret::SecretType type,
+                         DenseIntElementsAttr attr) -> std::optional<Type> {
+      return secret::SecretType::get(materializePermutationLayout(
+          getElementTypeOrSelf(type.getValueType()), attr,
+          getCiphertextSize()));
+    });
+    addConversion([this](RankedTensorType type,
+                         DenseIntElementsAttr attr) -> std::optional<Type> {
+      return materializePermutationLayout(getElementTypeOrSelf(type), attr,
+                                          getCiphertextSize());
+    });
   }
 
   int getCiphertextSize() const { return ciphertextSize; }
@@ -352,7 +363,7 @@ class ConvertAssignLayout
 
     // Check cache for existing assign layout function.
     Attribute layout = op.getLayout();
-    if (!isa<LayoutAttr>(layout)) {
+    if (!isa<LayoutAttr>(layout) && !isa<DenseIntElementsAttr>(layout)) {
       return failure();
     }
     Type inputType = input.getType();

lib/Transforms/ConvertToCiphertextSemantics/TypeConversion.cpp

@@ -41,5 +41,13 @@ Type materializeScalarLayout(Type type, LayoutAttr attr, int ciphertextSize) {
   return RankedTensorType::get({1, ciphertextSize}, type);
 }
 
+Type materializePermutationLayout(Type elementType,
+                                  DenseIntElementsAttr permutation,
+                                  int ciphertextSize) {
+  // TODO(#2666): Extend to a more general case where src_ct and dst_ct != 0
+  // src_ct and dst_ct are always 0; output is always a single ciphertext.
+  return RankedTensorType::get({1, ciphertextSize}, elementType);
+}
+
 }  // namespace heir
 }  // namespace mlir

lib/Transforms/ConvertToCiphertextSemantics/TypeConversion.h

@@ -18,6 +18,13 @@ Type materializeLayout(Type dataType, tensor_ext::LayoutAttr attr,
 Type materializeScalarLayout(Type type, tensor_ext::LayoutAttr attr,
                              int ciphertextSize);
 
+// Computes the ciphertext-semantic type for a permutation layout given as a
+// <N x 4 x i64> DenseIntElementsAttr of (src_ct, src_slot, dst_ct, dst_slot)
+// tuples. The number of ciphertexts is derived from the max dst_ct value.
+Type materializePermutationLayout(Type elementType,
+                                  DenseIntElementsAttr permutation,
+                                  int ciphertextSize);
+
 }  // namespace heir
 }  // namespace mlir
 

tests/Transforms/convert_to_ciphertext_semantics/assign_layout.mlir

@@ -64,3 +64,48 @@ module {
     return %0 : !secret.secret<tensor<32xi16>>
   }
 }
+
+// -----
+
+// CHECK: func.func private @_assign_layout_{{[0-9]+}}
+// CHECK-SAME: %[[ARG0:.*]]: tensor<1x32xi16>) -> tensor<1x32xi16>
+// CHECK-DAG: %[[c4:.*]] = arith.constant 4 : index
+// CHECK-DAG: %[[c3:.*]] = arith.constant 3 : index
+// CHECK-DAG: %[[c5:.*]] = arith.constant 5 : index
+// CHECK-DAG: %[[c2:.*]] = arith.constant 2 : index
+// CHECK-DAG: %[[c6:.*]] = arith.constant 6 : index
+// CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[c7:.*]] = arith.constant 7 : index
+// CHECK-DAG: %[[cst:.*]] = arith.constant dense<0> : tensor<1x32xi16>
+// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[EXT0:.*]] = tensor.extract %arg0[%[[c0]], %[[c0]]]
+// CHECK-DAG: %[[INS0:.*]] = tensor.insert %[[EXT0]] into %[[cst]][%[[c0]], %[[c7]]]
+// CHECK-DAG: %[[EXT1:.*]] = tensor.extract %arg0[%[[c0]], %[[c1]]]
+// CHECK-DAG: %[[INS1:.*]] = tensor.insert %[[EXT1]] into %[[INS0]][%[[c0]], %[[c6]]]
+// CHECK-DAG: %[[EXT2:.*]] = tensor.extract %arg0[%[[c0]], %[[c2]]]
+// CHECK-DAG: %[[INS2:.*]] = tensor.insert %[[EXT2]] into %[[INS1]][%[[c0]], %[[c5]]]
+// CHECK-DAG: %[[EXT3:.*]] = tensor.extract %arg0[%[[c0]], %[[c3]]]
+// CHECK-DAG: %[[INS3:.*]] = tensor.insert %[[EXT3]] into %[[INS2]][%[[c0]], %[[c4]]]
+// CHECK-DAG: %[[EXT4:.*]] = tensor.extract %arg0[%[[c0]], %[[c4]]]
+// CHECK-DAG: %[[INS4:.*]] = tensor.insert %[[EXT4]] into %[[INS3]][%[[c0]], %[[c3]]]
+// CHECK-DAG: %[[EXT5:.*]] = tensor.extract %arg0[%[[c0]], %[[c5]]]
+// CHECK-DAG: %[[INS5:.*]] = tensor.insert %[[EXT5]] into %[[INS4]][%[[c0]], %[[c2]]]
+// CHECK-DAG: %[[EXT6:.*]] = tensor.extract %arg0[%[[c0]], %[[c6]]]
+// CHECK-DAG: %[[INS6:.*]] = tensor.insert %[[EXT6]] into %[[INS5]][%[[c0]], %[[c1]]]
+// CHECK-DAG: %[[EXT7:.*]] = tensor.extract %arg0[%[[c0]], %[[c7]]]
+// CHECK-DAG: %[[INS7:.*]] = tensor.insert %[[EXT7]] into %[[INS6]][%[[c0]], %[[c0]]]
+
+// CHECK: @permutate_vector
+#layout = dense<[[0, 0, 0, 7], [0, 1, 0, 6], [0, 2, 0, 5], [0, 3, 0, 4], [0, 4, 0, 3], [0, 5, 0, 2], [0, 6, 0, 1], [0, 7, 0, 0]]> : tensor<8x4xi64>
+module {
+  func.func @permutate_vector() {
+    %cst = arith.constant dense<1> : tensor<1x32xi16>
+    // CHECK: %[[cst:.*]] = arith.constant dense<1> : tensor<1x32xi16>
+    // CHECK: func.call @_assign_layout_{{[0-9]+}}(%[[cst]])
+    %0 = secret.generic() {
+      %1 = tensor_ext.assign_layout %cst {layout = #layout, tensor_ext.layout = #layout} : tensor<1x32xi16>
+      secret.yield %1 : tensor<1x32xi16>
+    } -> (!secret.secret<tensor<1x32xi16>> {tensor_ext.layout = #layout})
+    return
+  }
+}