Add `tensor_ext.combine` and an associated pattern to replace collapse-insertion-chains
- Dominant language
- MLIR
- Stars
- 906
- Forks
- 171
- Avg merge
- 4d 12h
- Merged PRs (30d)
- 32
Description
From the review on https://github.com/google/heir/pull/471
> In HECO, "left-over" insert/extract ops instead get mapped to the `combine` op which lets you merge two tensors. The `combine` op essentially models the "masked addition" you sometimes get in batched FHE code: given a (plaintext) mask `m`, you do `a * m + b * !m`, letting you pick independently for each index whether it comes from `a` or `b`. While there's of course a canonicalization to handle the "full replacement" case that this pass handles, not all programs batch that nicely (of course, in the HECO examples, we focused on those that do :wink:). Because an `insert/extract` generally costs one (ptxt-ctxt) mask multiplication and one (much, much more expensive) rotation, mapping things to `combine` instead gives much better performance for those programs that don't batch perfectly. In addition, the `combine` op is pretty useful in itself to express batched FHE programs, so it'd be nice to add to `tensor_ext`. Unfortunately, the HECO implementation of `combine` has a rather messsy parser/printer and canonicalizer (though it gets the job done).
Contributor guide
Assessment
This issue has not been assessed yet.