google / google/heir

Add `tensor_ext.combine` and an associated pattern to replace collapse-insertion-chains

Open
#512 2 comments 0 reactions 0 assignees View on GitHub
dialect: tensor_ext enhancement newcomer project
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.