layout: Add support for permutation layout in unpackOp - #3157
layout: Add support for permutation layout in unpackOp#3157VedantParanjape wants to merge 1 commit into
Conversation
Permutation layout is supported by assign_layout op. The inverse of this op is unpack op. It is required by client decryption handles to decrypt the data and store it according to the layout attribute. This patch adds support for permutation layout to work with unpackOp.
j2kun
left a comment
j2kun
left a comment
There was a problem hiding this comment.
I think we can merge this as-is, but I think the attribute as a whole needs to be changed to avoid the special cases and confusing aspects of this PR.
| Value result = zeroCtxt.getResult(); | ||
|
|
||
| int64_t ctBound = (tensorType.getRank() == 2) ? tensorType.getDimSize(0) : 1; | ||
| int64_t srcCtBound = |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
I just realized this, I can rework this when I address the overall issue with the design of layout attribute.
| // Tests `tensor_ext.unpack` with a dense-permutation layout attribute | ||
|
|
||
| // ----------------------------------------------------------------------------- | ||
| // Case 1: single dense permutation with a rank-1 data-semantic target. |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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...
There was a problem hiding this comment.
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, ...)
There was a problem hiding this comment.
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.
2 participants
Permutation layout is supported by assign_layout op. The inverse of this op is unpack op. It is required by client decryption handles to decrypt the data and store it according to the layout attribute. This patch adds support for permutation layout to work with unpackOp.