Determine how to generically identify and handle plaintext-ciphertext ops in `insert-rotate`
- Dominant language
- MLIR
- Stars
- 906
- Forks
- 171
- Avg merge
- 4d 12h
- Merged PRs (30d)
- 32
Description
The gx-kernel HECO example ported from https://github.com/google/heir/pull/587 shows a shortcoming of the existing `insert-rotate` pass.
Specifically, it anchors on binary operations that consume the result of extracting values from (1D) tensors, but in the gx-kernel example one of those tensors is a 2D tensor containing the weights of a convolution kernel. We got around that by canonicalizing the constant weight matrix into individual arith.constant ops and then having special patterns that identify constants and splat them into tensors.
However, there are two analogous situations:
1. A not-quite-constant but still plaintext 1D tensor is combined with bona-fide secret tensors. By 'not quite constant' I mean that `canonicalize` won't split it up into individual `arith.constant` ops, but it is still semantically a cleartext. E.g., it could be a function input that is not annotated with `{secret.secret}`, which might be further obscured by having some arith ops locally process it before the extract op.
2. Non-tensor, but non-constant scalar values that need to be splatted to be combined via `arith.muli` or similar to a normal "ciphertext" tensor value. Again, a scalar non-secret function input that is processed inside the function by arith ops seems like a sufficiently general case of this.
Both of the examples above suggest to me that we should have some pre-insert-rotate analysis pass that annotates which SSA values are restricted by the "tensor_ext.rotate" paradigm and which are not. Call such values "restricted" and "unrestricted" for brevity. For unrestricted SSA values, we can freely use `tensor.extract` and other arbitrary tensor ops, and splat scalars whenever they need to be combined with a restricted value. The `insert-rotate` pass needn't concern itself with how these identifications were chosen, but can match on the annotations in its patterns.
Contributor guide
Assessment
This issue has not been assessed yet.