google / google/heir

Loop support: prefer `linalg` over generic `affine.for`

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

Description

HEIR has supported `affine.for`, but in the eyes of dataflow analysis, generic `affine.for` with _arbitrary body_ is hard to understand (understanding loop has always been one of the hard part of compiler). I propose using `linalg` to describe what we have now.

I further argue (not 100% sure) that any `affine.for` that _could_ be understood/accepted by FHE-specific analysis can always be rewritten in `linalg` with primitives like `elementwise`/`reduce`, as the task of the analysis is to understand what part of `affine.for`'s body is elementwise and how the loop-carried variable forms a reduction. The is somewhat similar to the map-reduce programming model.

### `affine.for` usage

Semantically there are two places we use `affine.for`

1. User input (data semantic), like roberts_cross/box_blur where the algorithm is naturally described in for loops
2. tensors of ciphertext (ciphertext semantic), like halevi-shoup matmul.

#### Data semantic

For data semantic, we already need to mitigate it in some way.

1. As we are already using full-loop-unroll and heir-vectorizer to handle these loops, further lowering into the FHE world should not see them
2. Alternatively, the layout packing pass is responsible for transforming the data semantic loop into ciphertext-semantic loop, then that pass could also rewrite it in `linalg`.

So the problem reduces to loop support in ciphertext semantic.

#### Common ciphertext-semantic patterns in current HEIR codebase

There are quite a lot of pattern like

```
affine.for tensor
tensor.extract one element
computation
tensor.insert into loop-carried variable
```

or
```
affine.for memref
memref.load one element
computation
memref.store into memory
```
that is actually elementwise op in input + reduction on the loop carried variable. Instead of letting the lowering emit such _fused_ loop and let analysis divide them apart, we can divide them apart early like

```
%res_tensor = linalg.elementwise_binary { body }
%res_scalar = linalg.reduce %res_tensor
```

#### Case study for the matmul kernel

The halevi-shoup lowering should lower from data-semantic linalg.matmul to cipher-text-semantic combinations of `linalg.some_op`.

The current lowering in main is like

```
%2 is the reduction variable (result)
%input0 is the input vector x
%arg3 is progressive rotation of x
%3:2 = affine.for %arg1 = 1 to 1024 iter_args(%arg2 = %2, %arg3 = %input0) -> (tensor<1x1024xf32>, tensor<1x1024xf32>) {
%4 = tensor_ext.rotate %arg3, %c1 : tensor<1x1024xf32>, index
%extracted_slice = tensor.extract_slice %cst_0[%arg1, 0] [1, 1024] [1, 1] : tensor<1024x1024xf32> to tensor<1x1024xf32>
%5 = arith.mulf %4, %extracted_slice : tensor<1x1024xf32>
%6 = arith.addf %arg2, %5 : tensor<1x1024xf32>
affine.yield %6, %4 : tensor<1x1024xf32>, tensor<1x1024xf32>
}
```

The elementwise part of it is clearly `mulf` (with tensor splatting), and reduction is `addf`, but the `rotate` part is quite hard to understand. In analysis we have to write an individual rule for it to know that it is progressive rotation, and NoiseAnalysis is quite unhappy with it, see https://github.com/google/heir/issues/1517#issuecomment-2699449481

Instead we could have

```
%rotated_tensor_of_input = linalg.elementwise { rotate } on %input_splatted
%mul_result = linalg.elementwise { mulf } on %weight and %rotated_tensor_of_input
%result = linalg.reduce %mul_result
```

And the `%rotated_tensor_of_input` can be produced in many ways like 8-4-2-1 or BSGS (32-16-8-4/3,2,1) way instead of rotate by each index / progressively rotate.

### Analysis requirement

There are these following types of analysis in HEIR FHE world that must properly handle loop to have correct lowering

1. Secretness Analysis (for whether it is secret or not)
2. Level analysis (for RNS level)
3. Dimension Analysis (for ciphertext size)
4. Scale Analysis (for CKKS/BGV scale management)
5. Noise Analysis (for correctly determining noise bound and parameter selection)

With `linalg` clearly stating the batch behavior instead of generic `affine.for`, the lives of these analyses could be much easier. Otherwise each of them need to separate the loop into different semantic part and understand it (quick reflection: how _could_ analysis understand the `rotate %c1` in halevi-shoup kernel above)

Another point that supports `linalg` is that, as `if` is hard to analyse, we already have convert-if-to-select so it becomes something like `.filter` in map-reduce language. And it can be expressed in `linalg.elementwise`.

### Discussion on additional cost

Separate them forms much bigger memory cost, as what map-reduce would do. For now all the partial results are stored instead of being reduced immediately.

We should seek optimization from the loop-transformation world like fusion/tiling, after the FHE lowering. I think this is quite standard way like semantic first and optimization later instead of writing an optimized `affine.for` early.

### Discussion on lowering

We eventually need to lower it. I think that is at the boundary of the backend where linalg is expressed in c++ for.

### Discussion on frontend

If such reasoning above is acceptable, we might encourage frontend user / kernel writer to use map-reduce-ish language / functional programming way to express things instead of for loop as the former would be much semantically clear.

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.