Loop support: prefer `linalg` over generic `affine.for`
- 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
Assessment
This issue has not been assessed yet.