Support for affine.for in RotationAnalysis
- Dominant language
- MLIR
- Stars
- 906
- Forks
- 171
- Avg merge
- 4d 12h
- Merged PRs (30d)
- 32
Description
In #2711 the new RotationAnalysis supports scf.for loops, which is what the rolled FHE kernels use for their implementation, but if there are performance issues it may be useful to switch to using affine.for.
In this case, the kernels would likely need to switch to using `affine.apply` for their intermediate index calculations, and then the analysis of the set of rotation indices could be accomplished by: (my rough sketch of notes below)
Starting from the rotation op in question, fetching all enclosing affine ops
```cpp
SmallVector enclosingAffineOps;
affine::getEnclosingAffineOps(rotationOp.getOperation(), enclosingAffineOps);
```
Constructing a set of constraints
```cpp
affine::FlatAffineValueConstraints domain;
for (Operation* affineOp : enclosingAffineOps) {
// affine.for: addAffineForOpDomain
// affine.parallel_for: addAffineParallelForOpDomain
// affine.if: addAffineIfOpDomain
}
```
Then get the AffineMap defining the operand in the inner-most loop nest
```cpp
AffineMap map = operand.getAffineMap();
// this is for affine.load, would it also work for a tensor-based op? Hopefully we would not have to switch to using memrefs...
```
Convert the map to a relation and compose.
```cpp
FlatAffineRelation rel(map);
rel.appendConstraints(domain);
```
The 'Value' dimensions are now constrained by the 'IV' dimensions. Project out the IVs to get the Value Set. This leaves you with the constraints on the materialized integers
```
rel.projectOut(0, numIVs);
```
Then [enumerate](https://github.com/google/heir/blob/3970331d06ef6ca5c90a60382c49c26b04fcb203/lib/Utils/Layout/Utils.cpp#L839) the relation to get all the actual values
Contributor guide
Assessment
This issue has not been assessed yet.