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

SmallVector<Operation*> enclosingAffineOps; affine::getEnclosingAffineOps(rotationOp.getOperation(), enclosingAffineOps);

Constructing a set of constraints

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

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.

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 the relation to get all the actual values