google / google/heir

Slight optimization to choice of "anyRangePoint" in tensor.extract kernel

Open
#2,216 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
MLIR
Stars
906
Forks
171
Avg merge
4d 12h
Merged PRs (30d)
32

Description

In https://github.com/google/heir/pull/2205, the tensor.extract kernel uses anyRangePoint to decide what to mask from the ciphertext, and then uses convert_layout to ensure the resulting scalar is packed according to the result layout chosen by `layout-optimization`.

Suppose we have a case where the ciphertext contains the value to extract in two separate slots, and one aligns exactly with the desired output layout:

```
ct = [x, x, 7, x, x, 7, x]
desired_result = [0, 0, 7, 0, 0, 0, 0]
```

In this case the `anyRangePoint` may end up masking for the second 7, which then incurs a single extra rotation to align it to the right final position. Instead, we could iterate over the desired result layout's range points, find which slots are already set in corresponding range points of the ciphertext tensor, and then mask all of the matching slots. This will at best remove all rotations, and at worst get the maximum number of matching slots pre-populated, which should make the resulting convert_layout cheaper.

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.