Adds ElementAttr to tensor_ext.layout#2653
Conversation
|
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. |
asraa
left a comment
There was a problem hiding this comment.
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
|
@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?) |
I have added a test case, but I need to fix the verifier before this testcase actually passes. |
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. |
asraa
left a comment
There was a problem hiding this comment.
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!
j2kun
left a comment
There was a problem hiding this comment.
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> |
There was a problem hiding this comment.
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 |
There was a problem hiding this comment.
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)
|
Ah, you may also need to run |
9aed109 to
53dd11d
Compare
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.
53dd11d to
a1df9b9
Compare
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.