Merge pull request #1906 from ZenithalHourlyRate:mgmt-init-multi-use

PiperOrigin-RevId: 771058735

Author: copybara-github

lib/Dialect/Mgmt/IR/MgmtOps.cpp

@@ -7,6 +7,10 @@ namespace mlir {
 namespace heir {
 namespace mgmt {
 
+//===----------------------------------------------------------------------===//
+// Canonicalization Patterns
+//===----------------------------------------------------------------------===//
+
 // Kept inside a namespace because it generates a function called
 // populateWithGenerated, which can conflict with other generated patterns.
 #include "lib/Dialect/Mgmt/IR/MgmtCanonicalization.cpp.inc"
@@ -28,6 +32,16 @@ void AdjustScaleOp::getCanonicalizationPatterns(RewritePatternSet &results,
   results.add<ModReduceAfterAdjustScale>(context);
 }
 
+//===----------------------------------------------------------------------===//
+// Utils
+//===----------------------------------------------------------------------===//
+
+void cleanupInitOp(Operation *top) {
+  top->walk([&](mgmt::InitOp initOp) {
+    if (initOp->use_empty()) initOp.erase();
+  });
+}
+
 }  // namespace mgmt
 }  // namespace heir
 }  // namespace mlir

lib/Dialect/Mgmt/IR/MgmtOps.h

@@ -12,4 +12,15 @@
 #define GET_OP_CLASSES
 #include "lib/Dialect/Mgmt/IR/MgmtOps.h.inc"
 
+namespace mlir {
+namespace heir {
+namespace mgmt {
+
+/// Remove all unused mgmt.init ops from the top operation.
+void cleanupInitOp(Operation *top);
+
+}  // namespace mgmt
+}  // namespace heir
+}  // namespace mlir
+
 #endif  // LIB_DIALECT_MGMT_IR_MGMTOPS_H_

lib/Dialect/Mgmt/IR/MgmtOps.td

@@ -143,7 +143,9 @@ def Mgmt_AdjustScaleOp : Mgmt_Op<"adjust_scale"> {
   let hasCanonicalizer = 1;
 }
 
-def Mgmt_InitOp : Mgmt_Op<"init"> {
+def Mgmt_InitOp : Mgmt_Op<"init",
+    [MemoryEffects<[MemWrite]>, ElementwiseMappable, SameOperandsAndResultType, ConditionallySpeculatable]> {
+
   let summary = "Init the plaintext with mgmt attributes";
 
   let description = [{
@@ -161,13 +163,30 @@ def Mgmt_InitOp : Mgmt_Op<"init"> {
 
     To address the problem, for each _use_ of the plaintext, we insert an `mgmt.init`
     operation to initialize the plaintext with `mgmt` attributes.
+
+    Technical reasons for registering memory effects:
+
+    Register a (bogus) memory effect to prevent CSE from merging this op.
+    Two mgmt.init ops could be seen as equivalent only if they have the same
+    MgmtAttr with *level/dimension/scale* annotated, otherwise we could not
+    judge whether they are equivalent or not. In practice, we create the op first
+    and only in later analyses we know whether they are equivalent or not.
+
+    ConditionallySpeculatable is for isSpeculatable check in hoisting canonicalization.
   }];
 
   let arguments = (ins
     AnyType:$input
   );
   let results = (outs AnyType:$output);
   let assemblyFormat = "operands attr-dict `:` type($output)";
+
+  let extraClassDeclaration = [{
+    /// Interface method for ConditionallySpeculatable.
+    Speculation::Speculatability getSpeculatability() {
+      return Speculation::Speculatable;
+    }
+  }];
 }
 
 #endif  // LIB_DIALECT_MGMT_IR_MGMTOPS_TD_

lib/Dialect/Secret/Conversions/SecretToBGV/SecretToBGV.cpp

@@ -274,6 +274,7 @@ struct SecretToBGV : public impl::SecretToBGVBase<SecretToBGV> {
     }
 
     clearAttrs(getOperation(), mgmt::MgmtDialect::kArgMgmtAttrName);
+    mgmt::cleanupInitOp(getOperation());
   }
 };
 

lib/Dialect/Secret/Conversions/SecretToCKKS/SecretToCKKS.cpp

@@ -407,6 +407,7 @@ struct SecretToCKKS : public impl::SecretToCKKSBase<SecretToCKKS> {
     }
 
     clearAttrs(getOperation(), mgmt::MgmtDialect::kArgMgmtAttrName);
+    mgmt::cleanupInitOp(getOperation());
   }
 };
 

lib/Transforms/SecretInsertMgmt/SecretInsertMgmtBFV.cpp

@@ -93,9 +93,9 @@ struct SecretInsertMgmtBFV
     (void)walkAndApplyPatterns(getOperation(), std::move(patternsRelinearize));
 
     auto level = maxMulDepth;
-    // call Canonicalizer here because mgmt.init ops need to be moved out of the
-    // secret.generic.
-    // annotate mgmt attribute with all levels set to mulDepth
+    // 1. Canonicalizer moves mgmt::InitOp out of secret.generic.
+    // 2. AnnotateMgmt will merge level and dimension into MgmtAttr, for further
+    //   lowering. For B/FV, all levels should be set to mulDepth.
     OpPassManager pipeline("builtin.module");
     pipeline.addPass(createCanonicalizerPass());
     mgmt::AnnotateMgmtOptions annotateMgmtOptions;

lib/Transforms/SecretInsertMgmt/SecretInsertMgmtBGV.cpp

@@ -126,12 +126,12 @@ struct SecretInsertMgmtBGV
       (void)walkAndApplyPatterns(getOperation(), std::move(patternsMulDepth));
     }
 
-    // call Canonicalizer here because mgmt ops need to be ordered
-    // call CSE here because there may be redundant mod reduce
-    // one Value may get mod reduced multiple times in
-    // multiple Uses
-    //
-    // also run annotate-mgmt for lowering
+    // 1. Canonicalizer reorders mgmt ops like Rescale/LevelReduce/AdjustScale.
+    //    This is important for AnnotateMgmt.
+    //    Canonicalizer also moves mgmt::InitOp out of secret.generic.
+    // 2. CSE removes redundant mgmt::ModReduceOp.
+    // 3. AnnotateMgmt will merge level and dimension into MgmtAttr, for further
+    //   lowering.
     OpPassManager pipeline("builtin.module");
     pipeline.addPass(createCanonicalizerPass());
     pipeline.addPass(createCSEPass());

lib/Transforms/SecretInsertMgmt/SecretInsertMgmtCKKS.cpp

@@ -226,12 +226,12 @@ struct SecretInsertMgmtCKKS
       (void)walkAndApplyPatterns(getOperation(), std::move(patternsMulDepth));
     }
 
-    // call Canonicalizer here because mgmt ops need to be ordered
-    // call CSE here because there may be redundant mod reduce
-    // one Value may get mod reduced multiple times in
-    // multiple Uses
-    //
-    // also run annotate-mgmt for lowering
+    // 1. Canonicalizer reorders mgmt ops like Rescale/LevelReduce/AdjustScale.
+    //    This is important for AnnotateMgmt.
+    //    Canonicalizer also moves mgmt::InitOp out of secret.generic.
+    // 2. CSE removes redundant mgmt::ModReduceOp.
+    // 3. AnnotateMgmt will merge level and dimension into MgmtAttr, for further
+    //   lowering.
     OpPassManager pipeline("builtin.module");
     pipeline.addPass(createCanonicalizerPass());
     pipeline.addPass(createCSEPass());

tests/Transforms/secret_insert_mgmt/bgv/init.mlir

@@ -52,4 +52,48 @@ module {
     } -> !secret.secret<i32>
     return %0 : !secret.secret<i32>
   }
+
+  // CHECK: @pt_multiple_uses
+  // CHECK-SAME: (%[[arg0:.*]]: !secret.secret<i16> {mgmt.mgmt = #mgmt.mgmt<level = 2>},
+  // CHECK-SAME: %[[arg1:.*]]: i16)
+  func.func @pt_multiple_uses(%arg0: !secret.secret<i16>, %arg1: i16) -> (!secret.secret<i16>) {
+    // CHECK: %[[v0:.*]] = mgmt.init %[[arg1]] {mgmt.mgmt = #mgmt.mgmt<level = 2>} : i16
+    // CHECK: %[[v1:.*]] = mgmt.init %[[arg1]] {mgmt.mgmt = #mgmt.mgmt<level = 1>} : i16
+    %0 = secret.generic(%arg0: !secret.secret<i16>) {
+    ^body(%input0: i16):
+      // CHECK: arith.addi
+      // CHECK-SAME: %[[v0]]
+      %1 = arith.addi %input0, %arg1 : i16
+      // CHECK: arith.muli
+      %2 = arith.muli %1, %1 : i16
+      // CHECK: arith.muli
+      // CHECK-SAME: %[[v1]]
+      %3 = arith.muli %2, %arg1 : i16
+      secret.yield %3 : i16
+    } -> !secret.secret<i16>
+    return %0 : !secret.secret<i16>
+  }
+
+  // CHECK: @pt_multiple_uses_2
+  // CHECK-SAME: (%[[arg0:.*]]: !secret.secret<i16> {mgmt.mgmt = #mgmt.mgmt<level = 2>},
+  // CHECK-SAME: %[[arg1:.*]]: i16)
+  func.func @pt_multiple_uses_2(%arg0: !secret.secret<i16>, %arg1: i16) -> (!secret.secret<i16>) {
+    // CHECK: %[[v0:.*]] = mgmt.init %[[arg1]] {mgmt.mgmt = #mgmt.mgmt<level = 2>} : i16
+    // Note: these two mgmt.init should not merge, as later optimization like lazy relinearization
+    // or populate-scale will make them different.
+    // CHECK: %[[v1:.*]] = mgmt.init %[[arg1]] {mgmt.mgmt = #mgmt.mgmt<level = 2>} : i16
+    %0 = secret.generic(%arg0: !secret.secret<i16>) {
+    ^body(%input0: i16):
+      // CHECK: arith.addi
+      // CHECK-SAME: %[[v0]]
+      %1 = arith.addi %input0, %arg1 : i16
+      // CHECK: arith.muli
+      %2 = arith.muli %1, %1 : i16
+      // CHECK: arith.addi
+      // CHECK-SAME: %[[v1]]
+      %3 = arith.addi %2, %arg1 : i16
+      secret.yield %3 : i16
+    } -> !secret.secret<i16>
+    return %0 : !secret.secret<i16>
+  }
 }