RotationAnalysis: model nested affine.for in arithmetic DAGs - #3354
Open
AlexanderViand wants to merge 1 commit into
Open
RotationAnalysis: model nested affine.for in arithmetic DAGs#3354AlexanderViand wants to merge 1 commit into
AlexanderViand wants to merge 1 commit into
Conversation
RotationAnalysis rebuilds an enclosing scf.for as an ArithmeticDag and enumerates the shifts its rotations use, so configure-crypto-context can generate exactly the Galois keys a program needs. DagBuilder covered scf.for, scf.if and the tensor ops, but not an affine.for nested in that reconstructed loop. An unhandled nested op is resolved as an opaque variable, severing everything upstream of it from the DAG root. The analysis still reports success, but the rotations inside the affine loop disappear. On tcn, rotate(ct, %i + 64) over %i = 1..5 therefore omitted the key for rotation 65 and failed at runtime. Teach DagBuilder to model constant-bound affine.for and affine.yield just like their scf counterparts. The lit test exercises both the induction variable and induction-variable-plus-constant forms inside an affine.for nested in scf.for; the unit test covers DAG evaluation directly. This does not handle a rotation whose enclosing loop is a bare affine.for. RotationAnalysis still rejects that case under TODO(#2712); this change only prevents nested affine loops from severing an scf.for DAG.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
RotationAnalysis rebuilds an enclosing scf.for as an ArithmeticDag and enumerates the shifts its rotations use, so configure-crypto-context can generate exactly the Galois keys a program needs. DagBuilder covered scf.for, scf.if and the tensor ops, but not an affine.for nested in that reconstructed loop. The analysis still reports success, but the rotations inside the affine loop disappear.
This teaches DagBuilder to model constant-bound affine.for and affine.yield just like their scf counterparts. The lit test exercises both the induction variable and induction-variable-plus-constant forms inside an affine.for nested in scf.for; the unit test covers DAG evaluation directly.
(If you can't already tell from the language - this one is straight from the 🤖. It checks out to me and works for us, but I'm not sure if there are subtleties about either the RotationAnalysis or the DagBuilder I'm missing)