Handling Rotom output in HEIR
- Dominant language
- MLIR
- Stars
- 906
- Forks
- 171
- Avg merge
- 4d 12h
- Merged PRs (30d)
- 32
Description
@edwjchen @j2kun and I were discussing the Rotom integration https://github.com/google/heir/issues/2035 into HEIR. After the Rotom pass, we need to represent the assigned layouts and kernels in HEIR in order for HEIR to materialize the ciphertext types and kernel operations (in convert-to-ciphertext-semantics).
In summary we noted a few things
* We can represent the Rotom kernel as a string attribute consisting of the function name that Rotom can create that contains the kernel operations. Rotom's kernels consist of a series of layout transformations and component functions (e.g. accumulate, multiply) to apply to the inputs. The arguments of this function should be the materialized ciphertext types. So the Rotom output of a tensor operation would look like
```
func.func @matvec_kernel_123(%arg0, %arg1)
func.func @main {
...
linalg.matvec ins(%0, %1) outs(%2) { kernel = "matvec_kernel_123" }
...
}
```
* Rotom can utilize `tensor_ext.convert_layout` to let HEIR handle layout conversions between kernels, this would require Rotom to convert the Rotom layout to an ISL string attribute. HEIR uses vos-vos-erkin shiftnetworks to lower this operation in convert-to-ciphertext-semantics
* If Rotom needs to apply/assign a layout to a cleartext to create a plaintext, then it should use `tensor_ext.assign_layout`. HEIR later materializes this as an assignment in a nested for loop. I propose that we extend `tensor_ext.assign_layout` to also allow for a rotom layout attribute rather than just an isl layout attribute. Since Rotom's layouts are already [representable](https://github.com/cmu-cryptosystems/Rotom/blob/9cc5744b0c62e81dde842bd4814777a4ba8ae595/util/layout_util.py#L487) as for loops, I don't think it makes sense to generalize to ISL (Rotom layout -> ISL -> for loops), and risk the time it might take for ISL to convert back to a nested for loop. But this can also be determined if we stress test the ISL -> for loop path. (I think it would reduce down to the problem of ISL -> Rotom which seems hard to find).
Contributor guide
Assessment
This issue has not been assessed yet.