Skip to content

layout: Add support for permutation layout in unpackOp #3157

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

Open
VedantParanjape wants to merge 1 commit into
base: main
Choose a base branch
from
Open

layout: Add support for permutation layout in unpackOp #3157

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
108 changes: 70 additions & 38 deletions lib/Transforms/ConvertToCiphertextSemantics/AssignLayout.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,7 +132,7 @@ static FailureOr<Value> implementUnpackOpStep(
return loop.value().getResults()[0];
}

static FailureOr<Value> implementAssignLayoutPermutation(
static FailureOr<Value> implementLayoutPermutation(
Value input, DenseIntElementsAttr permutation, Type targetTypeTy,
ImplicitLocOpBuilder& builder,
const std::function<void(Operation*)>& createdOpCallback) {
Expand All @@ -150,8 +150,11 @@ static FailureOr<Value> implementAssignLayoutPermutation(
createdOpCallback(zeroCtxt);
Value result = zeroCtxt.getResult();

int64_t ctBound = (tensorType.getRank() == 2) ? tensorType.getDimSize(0) : 1;
int64_t srcCtBound =

@j2kun j2kun Jul 2, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The naming is a little strange here because only one of the source and destination are semantically "ct, slot," depending on the caller's context. Moreover, the non-ct-slot side is restricted (somewhat arbitrarily) to be dimension 2.

Maybe don't name the variables "ct, slot", and then document the function with a code comment for the dim-2 restriction and for the semantic ambiguity of which one is the data and which is ciphertext.

@VedantParanjape VedantParanjape Jul 2, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I just realized this, I can rework this when I address the overall issue with the design of layout attribute.

(tensorType.getRank() == 2) ? tensorType.getDimSize(0) : 1;
int64_t srcSlotBound = tensorType.getDimSize(tensorType.getRank() - 1);
int64_t dstCtBound =
(targetType.getRank() == 2) ? targetType.getDimSize(0) : 1;
int64_t dstSlotBound = targetType.getDimSize(targetType.getRank() - 1);

for (auto it = permutation.value_begin<APInt>();
Expand All @@ -161,14 +164,15 @@ static FailureOr<Value> implementAssignLayoutPermutation(
int64_t dstCt = (*it++).getSExtValue();
int64_t dstSlot = (*it++).getSExtValue();

if (srcCt >= ctBound || srcSlot >= srcSlotBound || dstCt >= ctBound ||
if (srcCt >= srcCtBound || srcSlot >= srcSlotBound || dstCt >= dstCtBound ||
dstSlot >= dstSlotBound) {
return builder.emitError()
<< "Permutation index out of bounds: " << "src_ct=" << srcCt
<< ", src_slot=" << srcSlot << " (input bounds: ct < " << ctBound
<< ", slot < " << srcSlotBound << "); " << "dst_ct=" << dstCt
<< ", dst_slot=" << dstSlot << " (target bounds: ct < " << ctBound
<< ", slot < " << dstSlotBound << ")";
<< ", src_slot=" << srcSlot << " (input bounds: ct < "
<< srcCtBound << ", slot < " << srcSlotBound << "); "
<< "dst_ct=" << dstCt << ", dst_slot=" << dstSlot
<< " (target bounds: ct < " << dstCtBound << ", slot < "
<< dstSlotBound << ")";
}

SmallVector<Value> extractIndices;
Expand All @@ -185,13 +189,17 @@ static FailureOr<Value> implementAssignLayoutPermutation(
auto extracted = tensor::ExtractOp::create(builder, input, extractIndices);
createdOpCallback(extracted);

auto dstCtIdxOp = arith::ConstantIndexOp::create(builder, dstCt);
createdOpCallback(dstCtIdxOp);
SmallVector<Value> insertIndices;
if (targetType.getRank() == 2) {
auto dstCtIdxOp = arith::ConstantIndexOp::create(builder, dstCt);
createdOpCallback(dstCtIdxOp);
insertIndices.push_back(dstCtIdxOp.getResult());
}
auto dstSlotIdxOp = arith::ConstantIndexOp::create(builder, dstSlot);
createdOpCallback(dstSlotIdxOp);
auto insertOp = tensor::InsertOp::create(
builder, extracted.getResult(), result,
ValueRange{dstCtIdxOp.getResult(), dstSlotIdxOp.getResult()});
insertIndices.push_back(dstSlotIdxOp.getResult());
auto insertOp = tensor::InsertOp::create(builder, extracted.getResult(),
result, insertIndices);
createdOpCallback(insertOp);

result = insertOp.getResult();
Expand Down Expand Up @@ -431,8 +439,8 @@ FailureOr<Value> implementAssignLayout(
dyn_cast<DenseIntElementsAttr>(layout)) {
Type targetType = materializePermutationLayout(input.getType(), elementAttr,
ciphertextSize);
return implementAssignLayoutPermutation(input, elementAttr, targetType,
builder, createdOpCallback);
return implementLayoutPermutation(input, elementAttr, targetType, builder,
createdOpCallback);
}
return builder.emitError() << "Unsupported layout attribute type: " << layout;
}
Expand All @@ -447,39 +455,63 @@ FailureOr<Value> implementUnpackOp(
createdOpCallback);
}

// For Dense permutation layout we reuse implementLayoutPermutation. The
// mechanical action is identical (extract from input[src_ct, src_slot],
// insert into result[dst_ct, dst_slot]); only the semantic direction
// differs.
if (auto elementAttr = dyn_cast<DenseIntElementsAttr>(op.getLayout())) {
return implementLayoutPermutation(op.getValue(), elementAttr,
op.getResult().getType(), builder,
createdOpCallback);
}

if (auto arrayAttr = dyn_cast<ArrayAttr>(op.getLayout())) {
Value currentInput = op.getValue();
Type finalTargetType = op.getResult().getType();

for (int i = arrayAttr.size() - 1; i >= 0; --i) {
auto layoutAttr = cast<LayoutAttr>(arrayAttr[i]);
Type stepTargetType;
if (i == 0) {
stepTargetType = finalTargetType;
} else {
presburger::IntegerRelation rel = layoutAttr.getIntegerRelation();
unsigned numDomainVars = rel.getNumDomainVars();
SmallVector<int64_t> domainShape;
for (unsigned d = 0; d < numDomainVars; ++d) {
auto ub = rel.getConstantBound64(
presburger::BoundType::UB,
rel.getVarKindOffset(presburger::VarKind::Domain) + d);
if (!ub.has_value()) {
return builder.emitError()
<< "Unbounded domain variable in relation";
Attribute member = arrayAttr[i];

// Dense-permutation member: dispatch to the permutation helper.
if (auto denseAttr = dyn_cast<DenseIntElementsAttr>(member)) {
auto res =
implementLayoutPermutation(currentInput, denseAttr, finalTargetType,
builder, createdOpCallback);
if (failed(res)) {
return failure();
}
currentInput = res.value();
continue;
} else if (auto layoutAttr = dyn_cast<LayoutAttr>(member)) {
Type stepTargetType;
if (i == 0) {
stepTargetType = finalTargetType;
} else {
presburger::IntegerRelation rel = layoutAttr.getIntegerRelation();
unsigned numDomainVars = rel.getNumDomainVars();
SmallVector<int64_t> domainShape;
for (unsigned d = 0; d < numDomainVars; ++d) {
auto ub = rel.getConstantBound64(
presburger::BoundType::UB,
rel.getVarKindOffset(presburger::VarKind::Domain) + d);
if (!ub.has_value()) {
return builder.emitError()
<< "Unbounded domain variable in relation";
}
domainShape.push_back(ub.value() + 1);
}
domainShape.push_back(ub.value() + 1);
Type elementType = getElementTypeOrSelf(finalTargetType);
stepTargetType = RankedTensorType::get(domainShape, elementType);
}
Type elementType = getElementTypeOrSelf(finalTargetType);
stepTargetType = RankedTensorType::get(domainShape, elementType);
}

auto res = implementUnpackOpStep(currentInput, layoutAttr, stepTargetType,
builder, createdOpCallback);
if (failed(res)) {
return failure();
auto res =
implementUnpackOpStep(currentInput, layoutAttr, stepTargetType,
builder, createdOpCallback);
if (failed(res)) {
return failure();
}
currentInput = res.value();
}
currentInput = res.value();
}
return currentInput;
}
Expand Down
64 changes: 64 additions & 0 deletions tests/Transforms/add_client_interface/dense_unpack.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
// RUN: heir-opt --lower-unpack --canonicalize %s | FileCheck %s

// Tests `tensor_ext.unpack` with a dense-permutation layout attribute

// -----------------------------------------------------------------------------
// Case 1: single dense permutation with a rank-1 data-semantic target.

@j2kun j2kun Jul 2, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the design of this explicit layout attribute is insufficient for what you're trying to use it for, and this PR is showing the strain.

The original permutation attribute seems incorrect in retrospect:

  // A list of tuples [a, b, c, d] representing an explicit map (ct, slot) ->
  // (ct, slot) defined by f(a, b) = (c, d).

In fact, it's not a mapping (ct, slot) -> (ct, slot), because one side is data semantic (and restricted to rank 2) and the other is ciphertext semantic.

And then here, one must further hack around it by having the permutation attribute use rank 2 indices for the data semantic side, even though it has semantic rank 1.

Probably what it should be is

  // A list of tuples [a, b, ..., c, d] representing an explicit map 
  // (d0, d1, ...) -> (ct, slot)
  // defined by f(a, b, ...) = (c, d)

With a verifier to assert that when this is attached to an op, the ranks agree (the data semantic tensor involved has rank denseEltsAttr.getType().getShape()[1] - 2 and numElements = denseEltsAttr.getType().getShape()[0]).

Then the layout attribute would naturally support any rank of data-semantic tensor.

@VedantParanjape VedantParanjape Jul 2, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes that seems like a more generic implementation. But the subtle issue with f(a, b, ...) = (c, d) is that it cannot be used on result attributes. How about I implement a n-dim to n-dim mapping and then let downstream passes in the ciphertext semantic verify its correctness.

@j2kun j2kun Jul 8, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

But the subtle issue with f(a, b, ...) = (c, d) is that it cannot be used on result attributes.

I don't see why not...

@VedantParanjape VedantParanjape Jul 9, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

But the subtle issue with f(a, b, ...) = (c, d) is that it cannot be used on result attributes.

I don't see why not...

because for a result attribute, we are going from ciphertext to a n-dim array, hence we would need f(a, b) = (c, d, e, ...)

@j2kun j2kun Jul 13, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The layout attribute is defined to always map from the data-semantic domain to the ciphertext-semantic domain. See for example https://github.com/google/heir/blob/main/tests/Transforms/lower_unpack/lower_unpack.mlir

Then unpack is supposed to be smart enough to know that it needs to invert the mapping.

//
// The permutation `[[0, 3, 0, 0], [0, 5, 0, 1], [0, 1, 0, 2]]` gathers slots
// 3, 5, 1 from a single ciphertext view (rank-2, one row) into positions
// 0, 1, 2 of a rank-1 result.

#dense_perm_single = dense<[[0, 3, 0, 0], [0, 5, 0, 1], [0, 1, 0, 2]]>
: tensor<3x4xi64>
#orig_single = #tensor_ext.original_type<
originalType = tensor<3xi32>, layout = #dense_perm_single>

// CHECK: @unpack_dense_single
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
// CHECK-DAG: %[[C5:.*]] = arith.constant 5 : index
// CHECK-DAG: %[[ZERO:.*]] = arith.constant dense<0> : tensor<3xi32>
// CHECK-DAG: %[[E0:.*]] = tensor.extract %arg0[%[[C0]], %[[C3]]] : tensor<1x8xi32>
// CHECK-DAG: tensor.insert %[[E0]] into %{{.*}}[%[[C0]]] : tensor<3xi32>
// CHECK-DAG: %[[E1:.*]] = tensor.extract %arg0[%[[C0]], %[[C5]]] : tensor<1x8xi32>
// CHECK-DAG: tensor.insert %[[E1]] into %{{.*}}[%[[C1]]] : tensor<3xi32>
// CHECK-DAG: %[[E2:.*]] = tensor.extract %arg0[%[[C0]], %[[C1]]] : tensor<1x8xi32>
// CHECK-DAG: tensor.insert %[[E2]] into %{{.*}}[%[[C2]]] : tensor<3xi32>
// CHECK: return %{{.*}} : tensor<3xi32>
func.func @unpack_dense_single(
%arg0: tensor<1x8xi32> {tensor_ext.original_type = #orig_single}
) -> tensor<3xi32> {
%0 = tensor_ext.unpack %arg0 {layout = #dense_perm_single}
: (tensor<1x8xi32>) -> tensor<3xi32>
return %0 : tensor<3xi32>
}

// -----------------------------------------------------------------------------
// Case 2: dense permutation with cross-ct source. The permutation
// `[[1, 2, 0, 0], [0, 4, 0, 1]]` says logical output 0 lives at (ct=1, slot=2)
// and logical output 1 lives at (ct=0, slot=4).

#dense_perm_cross = dense<[[1, 2, 0, 0], [0, 4, 0, 1]]> : tensor<2x4xi64>
#orig_cross = #tensor_ext.original_type<
originalType = tensor<2xi32>, layout = #dense_perm_cross>

// CHECK: @unpack_dense_cross_ct
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
// CHECK-DAG: %[[E0:.*]] = tensor.extract %arg0[%[[C1]], %[[C2]]] : tensor<2x8xi32>
// CHECK-DAG: tensor.insert %[[E0]] into %{{.*}}[%[[C0]]] : tensor<2xi32>
// CHECK-DAG: %[[E1:.*]] = tensor.extract %arg0[%[[C0]], %[[C4]]] : tensor<2x8xi32>
// CHECK-DAG: tensor.insert %[[E1]] into %{{.*}}[%[[C1]]] : tensor<2xi32>
// CHECK: return %{{.*}} : tensor<2xi32>
func.func @unpack_dense_cross_ct(
%arg0: tensor<2x8xi32> {tensor_ext.original_type = #orig_cross}
) -> tensor<2xi32> {
%0 = tensor_ext.unpack %arg0 {layout = #dense_perm_cross}
: (tensor<2x8xi32>) -> tensor<2xi32>
return %0 : tensor<2xi32>
}
Loading