From 820633b3de94c8ead3eb9a337341be1fde48fefa Mon Sep 17 00:00:00 2001 From: Asra Ali Date: Tue, 29 Jul 2025 07:53:33 -0700 Subject: [PATCH] cleanup: use helper methods from secretness analysis PiperOrigin-RevId: 788453260 --- .../SecretnessAnalysis/SecretnessAnalysis.cpp | 6 ++++++ .../SecretnessAnalysis/SecretnessAnalysis.h | 2 ++ .../ConvertIfToSelect/ConvertIfToSelect.cpp | 6 ++---- .../ConvertSecretExtractToStaticExtract.cpp | 14 +++----------- .../ConvertSecretInsertToStaticInsert.cpp | 18 ++++++------------ .../ConvertSecretWhileToStaticFor.cpp | 10 ++-------- 6 files changed, 21 insertions(+), 35 deletions(-) diff --git a/lib/Analysis/SecretnessAnalysis/SecretnessAnalysis.cpp b/lib/Analysis/SecretnessAnalysis/SecretnessAnalysis.cpp index 707fdaf91d..4521bab0a8 100644 --- a/lib/Analysis/SecretnessAnalysis/SecretnessAnalysis.cpp +++ b/lib/Analysis/SecretnessAnalysis/SecretnessAnalysis.cpp @@ -224,6 +224,12 @@ bool isSecret(Value value, DataFlowSolver *solver) { return isSecret(lattice); } +bool isInitialized(Value value, DataFlowSolver *solver) { + auto *lattice = solver->lookupState(value); + if (!lattice) return false; + return lattice->getValue().isInitialized(); +} + bool isSecret(const SecretnessLattice *lattice) { if (!lattice) { return false; diff --git a/lib/Analysis/SecretnessAnalysis/SecretnessAnalysis.h b/lib/Analysis/SecretnessAnalysis/SecretnessAnalysis.h index b1fc5ef31c..743b4dfe4e 100644 --- a/lib/Analysis/SecretnessAnalysis/SecretnessAnalysis.h +++ b/lib/Analysis/SecretnessAnalysis/SecretnessAnalysis.h @@ -226,6 +226,8 @@ void annotateSecretness(Operation *top, DataFlowSolver *solver, bool verbose); // analysis bool isSecret(Value value, DataFlowSolver *solver); +bool isInitialized(Value value, DataFlowSolver *solver); + bool isSecret(const SecretnessLattice *lattice); bool isSecret(ValueRange values, DataFlowSolver *solver); diff --git a/lib/Transforms/ConvertIfToSelect/ConvertIfToSelect.cpp b/lib/Transforms/ConvertIfToSelect/ConvertIfToSelect.cpp index 45485cc1fd..b0d7ab7843 100644 --- a/lib/Transforms/ConvertIfToSelect/ConvertIfToSelect.cpp +++ b/lib/Transforms/ConvertIfToSelect/ConvertIfToSelect.cpp @@ -38,12 +38,10 @@ struct IfToSelectConversion : OpRewritePattern { LogicalResult matchAndRewrite(scf::IfOp ifOp, PatternRewriter &rewriter) const override { - auto *lattice = solver->lookupState(ifOp.getOperand()); - Secretness secretness = lattice ? lattice->getValue() : Secretness(); - // Convert ops with "secret" and, conservatively, "unknown" (uninitialized) // conditions but skip conversion if the condition is known to be non-secret - if (secretness.isInitialized() && !secretness.getSecretness()) + auto ifOperand = ifOp.getOperand(); + if (isInitialized(ifOperand, solver) && !isSecret(ifOperand, solver)) return failure(); // Hoist instructions in the 'then' and 'else' regions diff --git a/lib/Transforms/ConvertSecretExtractToStaticExtract/ConvertSecretExtractToStaticExtract.cpp b/lib/Transforms/ConvertSecretExtractToStaticExtract/ConvertSecretExtractToStaticExtract.cpp index bab9d151a5..ae0792271f 100644 --- a/lib/Transforms/ConvertSecretExtractToStaticExtract/ConvertSecretExtractToStaticExtract.cpp +++ b/lib/Transforms/ConvertSecretExtractToStaticExtract/ConvertSecretExtractToStaticExtract.cpp @@ -49,25 +49,17 @@ struct SecretExtractToStaticExtractConversion auto index = extractOp.getIndices().front(); - auto *indexSecretnessLattice = - solver->lookupState(index); - - // Use secretness from lattice or, if no lattice found, set to unknown - auto indexSecretness = indexSecretnessLattice - ? indexSecretnessLattice->getValue() - : Secretness(); - // If lattice is set to unknown, // apply transformation anyway but emit a warning - if (!indexSecretness.isInitialized()) { + bool initialized = isInitialized(index, solver); + if (!initialized) { extractOp->emitWarning() << "Secretness for tensor.extract index is unknown. " "Conservatively, the transformation will be applied:"; } // If index is known to be public, no transformation is needed - if (indexSecretness.isInitialized() && !indexSecretness.getSecretness()) - return failure(); + if (initialized && !isSecret(index, solver)) return failure(); ImplicitLocOpBuilder builder(extractOp->getLoc(), rewriter); diff --git a/lib/Transforms/ConvertSecretInsertToStaticInsert/ConvertSecretInsertToStaticInsert.cpp b/lib/Transforms/ConvertSecretInsertToStaticInsert/ConvertSecretInsertToStaticInsert.cpp index 90e0027083..24349b905b 100644 --- a/lib/Transforms/ConvertSecretInsertToStaticInsert/ConvertSecretInsertToStaticInsert.cpp +++ b/lib/Transforms/ConvertSecretInsertToStaticInsert/ConvertSecretInsertToStaticInsert.cpp @@ -50,25 +50,19 @@ struct SecretInsertToStaticInsertConversion auto tensor = insertOp.getDest(); auto insertedValue = insertOp.getScalar(); - auto *indexSecretnessLattice = - solver->lookupState(index); - - // use secretness from lattice or, if no lattice found, set to unknown - auto indexSecretness = indexSecretnessLattice - ? indexSecretnessLattice->getValue() - : Secretness(); - - // If lattice is set to unknown, - // apply transformation anyway but emit a warning - if (!indexSecretness.isInitialized()) { + // If lattice is set to unknown, apply transformation anyway but emit a + // warning + bool initialized = isInitialized(index, solver); + if (!initialized) { insertOp->emitWarning() << "Secretness for tensor.insert index is unknown. " "Conservatively, the transformation will be applied:"; } // If index is known to be public, no transformation is needed - if (indexSecretness.isInitialized() && !indexSecretness.getSecretness()) + if (initialized && !isSecret(index, solver)) { return failure(); + } ImplicitLocOpBuilder builder(insertOp->getLoc(), rewriter); diff --git a/lib/Transforms/ConvertSecretWhileToStaticFor/ConvertSecretWhileToStaticFor.cpp b/lib/Transforms/ConvertSecretWhileToStaticFor/ConvertSecretWhileToStaticFor.cpp index 81fc254984..936146bc4d 100644 --- a/lib/Transforms/ConvertSecretWhileToStaticFor/ConvertSecretWhileToStaticFor.cpp +++ b/lib/Transforms/ConvertSecretWhileToStaticFor/ConvertSecretWhileToStaticFor.cpp @@ -63,21 +63,15 @@ struct SecretWhileToStaticForConversion : OpRewritePattern { counter++; } - auto *conditionSecretnessLattice = - solver->lookupState(whileCondition); - - if (!conditionSecretnessLattice) { + if (!isInitialized(whileCondition, solver)) { InFlightDiagnostic diag = whileOp.emitWarning() << "Secretness for scf.while condition has not been set"; return failure(); } - bool isConditionSecret = - conditionSecretnessLattice->getValue().getSecretness(); - // If condition is not secret, no transformation is needed - if (!isConditionSecret) { + if (!isSecret(whileCondition, solver)) { return failure(); }