google / google/heir

Write a linalg.reduce kernel using the relation-based layout

Open
#2,254 3 comments 0 reactions 0 assignees View on GitHub
good first issue
Dominant language
MLIR
Stars
906
Forks
171
Avg merge
4d 12h
Merged PRs (30d)
32

Description

With the new layout system, a linalg.reduce kernel is a bit more complicated, mainly due to the extra decisions possible with a multi-ciphertext packing.

In the old setup, we only allowed data to be packed into a single ciphertext, so the linalg.reduce kernel involved aligning the axis to be summed as a stride, and then doing a partial rotate-and-reduce to reduce them.

In the new setup, given an arbitrary input packing, we could utilize this same approach, or align slots across multiple ciphertexts, or a combination of both. I'm not sure what the right approach is (haven't thought about it too hard), but in general it may require repacking the ciphertext to make slots align.

What we should do is decide on some kernel options, add them to layout-optimization so it can handle the implied layout conversions required to make the kernel feasible, then implement the kernel chosen by the optimizer in ConvertToCiphertextSemantics for the assumed layout.

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.