google / google/heir

Support for affine.for in RotationAnalysis

Open
#2,712 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 #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

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.