Skip to content

Adds ElementAttr to tensor_ext.layout#2653

Merged
copybara-service[bot] merged 1 commit intogoogle:mainfrom
VedantParanjape:permutation
Feb 18, 2026
Merged

Adds ElementAttr to tensor_ext.layout#2653
copybara-service[bot] merged 1 commit intogoogle:mainfrom
VedantParanjape:permutation

Conversation

@VedantParanjape
Copy link
Copy Markdown
Collaborator

TensorExt.Layout could only be assigned a LayoutAttr, which is an quasi-affine relation. This adds support to handle arbitrary permutations of ctxt to vector slots.

@google-cla
Copy link
Copy Markdown

google-cla Bot commented Feb 13, 2026

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

Copy link
Copy Markdown
Collaborator

@asraa asraa left a comment

Choose a reason for hiding this comment

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

Hey @VedantParanjape ! Code looks really great, thanks for the contribution.

I was just wondering if you could add a few small tests with a permutation layout in tests/Transforms/convert_to_ciphertext_semantics

Comment thread lib/Transforms/ConvertToCiphertextSemantics/AssignLayout.cpp Outdated
@j2kun
Copy link
Copy Markdown
Collaborator

j2kun commented Feb 13, 2026

@VedantParanjape In some of the examples we discussed, the packing was not purely a permutation (there were replicated values among slots). I worry that a strict "permutation" does not suffice, when actually you want a mapping from slots to data values.

In the case where we do remapping from (ct, slot) -> (ct, slot), an AnyI64ElementsAttr suffices because we can group the values in groups of four (a, b, c, d) being interpreted as mapping f(a, b) = (c, d). Can you also do that in your case? (maybe yes because you're supporting flattened input tensors?)

@VedantParanjape
Copy link
Copy Markdown
Collaborator Author

Hey @VedantParanjape ! Code looks really great, thanks for the contribution.

I was just wondering if you could add a few small tests with a permutation layout in tests/Transforms/convert_to_ciphertext_semantics

I have added a test case, but I need to fix the verifier before this testcase actually passes.

@VedantParanjape
Copy link
Copy Markdown
Collaborator Author

VedantParanjape commented Feb 13, 2026

In the case where we do remapping from (ct, slot) -> (ct, slot), an AnyI64ElementsAttr suffices because we can group the values in groups of four (a, b, c, d) being interpreted as mapping f(a, b) = (c, d). Can you also do that in your case? (maybe yes because you're supporting flattened input tensors?)

yeah, since I am just supporting a flattened input tensor, the ct (a, b) will really just be 0 in all the cases, so we can just eliminate it?

Overall your design makes more sense, I will make the changes. I need to just duplicate MappingLike code.

Copy link
Copy Markdown
Collaborator

@asraa asraa left a comment

Choose a reason for hiding this comment

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

thanks! i suggest either generalizing to any src/ct index or forcing or validating that the ct indices are all zero

but code looks great as usual!

Comment thread lib/Dialect/TensorExt/IR/TensorExtOps.td Outdated
Comment thread lib/Transforms/ConvertToCiphertextSemantics/AssignLayout.cpp Outdated
Comment thread lib/Transforms/ConvertToCiphertextSemantics/AssignLayout.cpp Outdated
Comment thread lib/Transforms/ConvertToCiphertextSemantics/TypeConversion.cpp
@VedantParanjape
Copy link
Copy Markdown
Collaborator Author

@j2kun @asraa The current version passes the new test case locally. Could you please trigger the workflow job?

Copy link
Copy Markdown
Collaborator

@j2kun j2kun left a comment

Choose a reason for hiding this comment

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

LGTM! Added some minor comments, and when you're ready please squash all your commits and then ping me and I'll get it merged in.

// CHECK-DAG: %[[c7:.*]] = arith.constant 7 : i32
// CHECK-DAG: %[[cst:.*]] = arith.constant dense<0> : tensor<1x32xi16>
// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : i32
// CHECK-DAG: %[[EXT0:.*]] = tensor.extract %arg0[%[[c0]], %[[c0]]] : tensor<1x32xi16>
Copy link
Copy Markdown
Collaborator

@j2kun j2kun Feb 18, 2026

Choose a reason for hiding this comment

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

nit: in general I think it's safe to leave out things like the type suffix when it's otherwise clear from the context. Less noise in the test files to update later.

Type materializePermutationLayout(Type elementType,
DenseIntElementsAttr permutation,
int ciphertextSize) {
// TODO: Extend to a more general case where src_ct and dst_ct != 0
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.

File an issue and use the notation

// TODO(#<issue_number>): foo

And this action will trigger to put a link to the source line on the GH issue directly. (And notify if the issue is closed while the TODO is still in the codebase)

Same for the other TODO above (they can both link to the same issue)

@j2kun
Copy link
Copy Markdown
Collaborator

j2kun commented Feb 18, 2026

Ah, you may also need to run pre-commit run --all-files to ensure the formatters are applied.

@VedantParanjape VedantParanjape force-pushed the permutation branch 2 times, most recently from 9aed109 to 53dd11d Compare February 18, 2026 04:39
@j2kun j2kun added the pull_ready Indicates whether a PR is ready to pull. The copybara worker will import for internal testing label Feb 18, 2026
TensorExt.Layout could only be assigned a LayoutAttr, which is an
quasi-affine relation. This adds support to handle arbitrary
permutations of ctxt to vector slots.
@copybara-service copybara-service Bot merged commit 1d4155f into google:main Feb 18, 2026
7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pull_ready Indicates whether a PR is ready to pull. The copybara worker will import for internal testing

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants