extend handling shift network source conflicts

isabelfaulds · isabelfaulds · commit 61bc1803e329 · 2026-01-06T20:29:05.000-08:00
diff --git a/lib/Dialect/TensorExt/Transforms/ImplementShiftNetwork.cpp b/lib/Dialect/TensorExt/Transforms/ImplementShiftNetwork.cpp
@@ -44,7 +44,7 @@ namespace tensor_ext {
 #include "lib/Dialect/TensorExt/Transforms/Passes.h.inc"
 
 ShiftScheme VosVosErkinShiftNetworks::findShiftScheme(
-    const Mapping& mapping, ArrayRef<int64_t> shiftOrder) {
+    Mapping& mapping, ArrayRef<int64_t> shiftOrder) {
   CacheKey cacheKey = makeCacheKey(mapping, shiftOrder);
   if (schemeCache.count(cacheKey)) {
     return schemeCache[cacheKey];
@@ -91,6 +91,241 @@ ShiftScheme VosVosErkinShiftNetworks::findShiftScheme(
     }
   });
 
+  // Cleaning subroutine, identify if overlapping sources for targets
+  auto targetToSources = mapping.getTargetToSources();
+  bool hasMultiSourceTargets =
+      std::any_of(targetToSources.begin(), targetToSources.end(),
+                  [](const auto& kv) { return kv.second.size() > 1; });
+  LLVM_DEBUG({
+    llvm::dbgs() << "Has overlapping sources for targets - "
+                 << hasMultiSourceTargets << "\n";
+  });
+  if (hasMultiSourceTargets) {
+
+    // Create a full conflict graph
+    graph::UndirectedGraph<CtSlot> fullConflictGraph;
+    for (const auto& [target, sources] : targetToSources) {
+        for (const auto& ctmatch : sources) {
+          fullConflictGraph.addVertex(ctmatch.source);
+        }
+      }
+    SmallVector<int64_t> defaultTestShiftOrder = defaultShiftOrder(
+        mapping.getCiphertextSize() * mapping.getNumCiphertexts());
+
+    ShiftStrategy fullConflictStrategy =
+        evaluateShiftStrategy(mapping, defaultTestShiftOrder, true);
+
+    for (const auto& [roundNum, round] :
+         llvm::enumerate(fullConflictStrategy.getRounds())) {
+      if (roundNum == 0) continue;
+      auto posns = round.positions;
+      for (auto it1 = posns.begin(); it1 != posns.end(); ++it1) {
+        for (auto it2 = std::next(it1); it2 != posns.end(); ++it2) {
+          const SourceShift& ss1 = it1->first;
+          const SourceShift& ss2 = it2->first;
+          if (ss1.source != ss2.source && it1->second == it2->second) {
+            LLVM_DEBUG(llvm::dbgs()
+                       << "Round " << roundNum << ": collision between "
+                       << "{" << ss1.source.ct << "," << ss1.source.slot << "}"
+                       << " and " << "{" << ss2.source.ct << ","
+                       << ss2.source.slot << "}" << " at " << "{"
+                       << it1->second.ct << "," << it1->second.slot << "}\n");
+            fullConflictGraph.addEdge(ss1.source, ss2.source);
+          }
+        }
+      }
+    }
+
+    
+    struct SourceInfo {
+      int numConflicts;
+      int numInserts;
+    };
+
+    std::unordered_map<CtSlot, SourceInfo> sourceMap;
+
+    for (CtSlot vertex : fullConflictGraph.getVertices()) {
+      sourceMap[vertex] = SourceInfo{
+          static_cast<int>(fullConflictGraph.edgesIncidentTo(vertex).size()),
+          0};
+    }
+
+    LLVM_DEBUG({
+      llvm::dbgs() << "Full Conflict graph:\n";
+      for (CtSlot vertex : fullConflictGraph.getVertices()) {
+        llvm::dbgs() << vertex.ct << "," << vertex.slot << " <-> {";
+        for (CtSlot neighbor : fullConflictGraph.edgesIncidentTo(vertex)) {
+          llvm::dbgs() << neighbor.ct << "," << neighbor.slot << "; ";
+        }
+        llvm::dbgs() << "}\n";
+      }
+    });
+
+    // Create seeds for tests
+    static std::mt19937 masterRng(std::random_device{}());
+    static std::mt19937 rng(std::random_device{}());
+    std::uniform_int_distribution<std::uint32_t> seedDist;
+
+    // Store best mapping
+    int minNumRotations = std::numeric_limits<int>::max();
+    Mapping bestMapping = mapping;
+    graph::UndirectedGraph<CtSlot> bestConflictGraph;
+
+    // Copy the target to Sources for testing 
+    std::vector<std::pair<CtSlot, SmallVector<Mapping::CtMatch>*>> allTargets;
+    allTargets.reserve(targetToSources.size());
+    for (auto &entry : targetToSources) {
+      allTargets.emplace_back(entry.first, &entry.second);
+    }
+
+    // Run 50 tests
+    for (int i = 0; i < 50; ++i) {
+      // track source numberConflicts , inserts
+      std::unordered_map<CtSlot, SourceInfo> testSourceMap = sourceMap;
+
+      // create test mapping and conflictgraph
+      Mapping testMapping = mapping;
+      graph::UndirectedGraph<CtSlot> testConflictGraph;
+      
+      // shuffle targets and fill
+      std::uint32_t trialSeed = seedDist(masterRng);
+      rng.seed(trialSeed);
+      std::shuffle(allTargets.begin(), allTargets.end(), rng); 
+
+      for (auto &[slot, matchesPtr] : allTargets) {
+        SmallVector<Mapping::CtMatch> &matches = *matchesPtr;
+        std::vector<std::pair<CtSlot, SourceInfo*>> candidates;
+        candidates.reserve(matches.size());
+
+        // Get and sort sources by 1) inserts , 2) conflicts ascending
+        for (auto &match : matches) {
+           CtSlot src = match.source;
+           auto it = testSourceMap.find(src);
+            candidates.emplace_back(src, &it->second);
+        }
+        std::sort(candidates.begin(), candidates.end(),
+          [](const auto &a, const auto &b) {
+            const SourceInfo &sa = *a.second;
+            const SourceInfo &sb = *b.second;
+            if (sa.numInserts != sb.numInserts)
+              return sa.numInserts < sb.numInserts;
+            return sa.numConflicts < sb.numConflicts;
+        });
+
+        LLVM_DEBUG({
+          llvm::dbgs() << "Candidates after sort for target "
+                      << slot.ct << "," << slot.slot << ":\n";
+          for (const auto &entry : candidates) {
+            const CtSlot &src = entry.first;
+            const SourceInfo &info = *entry.second;
+            llvm::dbgs() << "  src=(" << src.ct << "," << src.slot
+                        << "), numInserts=" << info.numInserts
+                        << ", numConflicts=" << info.numConflicts << "\n";
+          }
+        });
+
+        // Pick any source with lowest at random and save choice
+        const SourceInfo &best = *candidates.front().second;
+        int minInserts   = best.numInserts;
+        int minConflicts = best.numConflicts;
+
+        std::vector<size_t> bestIdx;
+        bestIdx.reserve(candidates.size());
+
+        // any source with lowest
+        for (size_t i = 0; i < candidates.size(); ++i) {
+          const SourceInfo &info = *candidates[i].second;
+          if (info.numInserts == minInserts && info.numConflicts == minConflicts)
+            bestIdx.push_back(i);
+          else
+            break;
+        }
+
+        // pick and save choice
+        std::uniform_int_distribution<size_t> dist(0, bestIdx.size() - 1);
+        size_t chosenIndex = bestIdx[dist(rng)];
+
+        CtSlot chosen = candidates[chosenIndex].first;
+        SourceInfo &info = *candidates[chosenIndex].second;
+        info.numInserts += 1;
+        testMapping.add(chosen, slot);
+        testConflictGraph.addVertex(chosen);
+      }
+
+
+      // Create test strategy, fill edges
+      ShiftStrategy testStrategy =
+          evaluateShiftStrategy(testMapping, defaultTestShiftOrder);
+      for (const auto& [roundNum, round] :
+         llvm::enumerate(testStrategy.getRounds())) {
+        if (roundNum == 0) continue;
+        auto posns = round.positions;
+        for (auto it1 = posns.begin(); it1 != posns.end(); ++it1) {
+          for (auto it2 = std::next(it1); it2 != posns.end(); ++it2) {
+            const SourceShift& ss1 = it1->first;
+            const SourceShift& ss2 = it2->first;
+            if (ss1.source != ss2.source && it1->second == it2->second) {
+              LLVM_DEBUG(llvm::dbgs()
+                        << "Round " << roundNum << ": collision between "
+                        << "{" << ss1.source.ct << "," << ss1.source.slot << "}"
+                        << " and " << "{" << ss2.source.ct << ","
+                        << ss2.source.slot << "}" << " at " << "{"
+                        << it1->second.ct << "," << it1->second.slot << "}\n");
+              testConflictGraph.addEdge(ss1.source, ss2.source);
+            }
+          }
+        }
+      }
+
+      // Find number of rotations needed
+      graph::GreedyGraphColoring<CtSlot> testColorer;
+      std::unordered_map<CtSlot, int> testColoring =
+          testColorer.color(testConflictGraph);
+
+      SmallVector<RotationGroup> testResultRotationGroups;
+      testResultRotationGroups.reserve(5);
+
+      for (const auto &entry : testColoring) {
+        CtSlot source = entry.first;
+        int64_t color = entry.second;
+        if (color >= static_cast<int64_t>(testResultRotationGroups.size()))
+          testResultRotationGroups.resize(color + 1);
+        testResultRotationGroups[color].insert(source);
+      }
+
+      // Update best number rotations & mapping
+      int testNumberRotations = static_cast<int>(testResultRotationGroups.size());
+      if (testNumberRotations < minNumRotations) {
+        minNumRotations = testNumberRotations;
+        bestMapping = testMapping;
+        bestConflictGraph = testConflictGraph;
+      }
+    }
+
+    mapping.setTargetToSource(bestMapping.getTargetToSource());
+    conflictGraph = bestConflictGraph;
+
+    LLVM_DEBUG({
+      llvm::dbgs() << "Cleaning subroutine finished. "
+                  << "Best num rotations = " << minNumRotations << "\n";
+
+      llvm::dbgs() << "Best mapping targetToSource:\n";
+      for (const auto &[target, source] : bestMapping.getTargetToSource()) {
+        llvm::dbgs() << "  target=(" << target.ct << "," << target.slot << ")"
+                    << " <- source=(" << source.ct << "," << source.slot << ")\n";
+      }
+
+      llvm::dbgs() << "Conflict graph after cleaning:\n";
+      for (CtSlot vertex : conflictGraph.getVertices()) {
+        llvm::dbgs() << "  " << vertex.ct << "," << vertex.slot << " <-> {";
+        for (CtSlot neighbor : conflictGraph.edgesIncidentTo(vertex)) {
+          llvm::dbgs() << neighbor.ct << "," << neighbor.slot << "; ";
+        }
+        llvm::dbgs() << "}\n";
+    } 
+    });
+  }
+
   graph::GreedyGraphColoring<CtSlot> colorer;
   std::unordered_map<CtSlot, int> coloring = colorer.color(conflictGraph);
 
@@ -125,7 +360,7 @@ ShiftScheme VosVosErkinShiftNetworks::findShiftScheme(
 }
 
 ShiftScheme VosVosErkinShiftNetworks::findBestShiftScheme(
-    const Mapping& mapping, std::size_t randomSeed, unsigned randomTries) {
+    Mapping& mapping, std::size_t randomSeed, unsigned randomTries) {
   SmallVector<int64_t> initShiftOrder = defaultShiftOrder(
       mapping.getCiphertextSize() * mapping.getNumCiphertexts());
 
@@ -156,15 +391,19 @@ ShiftScheme VosVosErkinShiftNetworks::findBestShiftScheme(
 }
 
 ShiftStrategy VosVosErkinShiftNetworks::evaluateShiftStrategy(
-    const Mapping& mapping, ArrayRef<int64_t> shiftOrder) {
+    const Mapping& mapping, ArrayRef<int64_t> shiftOrder, bool useSources) {
+
   CacheKey cacheKey = makeCacheKey(mapping, shiftOrder);
   if (strategyCache.count(cacheKey)) {
     return strategyCache[cacheKey];
   }
+  LLVM_DEBUG(llvm::dbgs() << "Making shift strategy..." << "\n");
 
   ShiftStrategy strategy(mapping.getCiphertextSize(),
                          mapping.getNumCiphertexts(), shiftOrder);
-  strategy.evaluate(mapping);
+  LLVM_DEBUG(llvm::dbgs() << "Evaluating shift strategy..." << "\n");
+
+  strategy.evaluate(mapping, useSources);
   strategyCache[cacheKey] = strategy;
   return strategy;
 }
diff --git a/lib/Dialect/TensorExt/Transforms/ImplementShiftNetwork.h b/lib/Dialect/TensorExt/Transforms/ImplementShiftNetwork.h
@@ -41,18 +41,19 @@ class VosVosErkinShiftNetworks {
   // on further calls to avoid recomputing the shift network.
   //
   // The default shiftOrder is LSB to MSB, i.e. 1, 2, 4, 8, ...
-  ShiftScheme findShiftScheme(const Mapping& mapping,
+  ShiftScheme findShiftScheme(Mapping& mapping,
                               ArrayRef<int64_t> shiftOrder = {});
 
   // Like findShiftScheme but randomly draw from a uniform distribution over all
   // possible shift orders and use the one that results in the best network.
-  ShiftScheme findBestShiftScheme(const Mapping& mapping,
+  ShiftScheme findBestShiftScheme(Mapping& mapping,
                                   std::size_t randomSeed,
                                   unsigned randomTries = 100);
 
  private:
   ShiftStrategy evaluateShiftStrategy(const Mapping& mapping,
-                                      ArrayRef<int64_t> shiftOrder);
+                                      ArrayRef<int64_t> shiftOrder,
+                                      bool useSources = false);
 
   CacheKey makeCacheKey(const Mapping& mapping, ArrayRef<int64_t> shiftOrder);
 
diff --git a/lib/Dialect/TensorExt/Transforms/ShiftScheme.cpp b/lib/Dialect/TensorExt/Transforms/ShiftScheme.cpp
@@ -42,13 +42,40 @@ int64_t ShiftStrategy::getVirtualShift(const CtSlot& source,
   return normalizeShift(sourceIndex, targetIndex, virtualCiphertextSize);
 }
 
-void ShiftStrategy::evaluate(const Mapping& mapping) {
+void ShiftStrategy::evaluate(const Mapping& mapping, bool useSources) {
   // First compute the virtual shifts needed for each source slot
   SmallVector<SourceShift> sourceShifts;
-  sourceShifts.reserve(mapping.size());
-  for (const auto& [target, source] : mapping.getTargetToSource()) {
-    int64_t shift = getVirtualShift(source, target);
-    sourceShifts.push_back({source, shift});
+  if (!useSources) {
+    sourceShifts.reserve(mapping.size());
+    for (const auto& [target, source] : mapping.getTargetToSource()) {
+      int64_t shift = getVirtualShift(source, target);
+      sourceShifts.push_back({source, shift});
+    }
+  } else {
+    LLVM_DEBUG(llvm::dbgs() << "Evaluating with all targets \n");
+
+    auto allTargets = mapping.getTargetToSources();
+
+    int counter = 0;
+    llvm::DenseSet<SourceShift> uniqueShifts;
+    for (const auto& kv : allTargets) {
+      const CtSlot& target = kv.first;
+      const auto& sources = kv.second;
+      for (const auto& ctmatch : sources) {
+        int64_t shift = getVirtualShift(ctmatch.source, target);
+        SourceShift ss{ctmatch.source, shift};
+
+    // Ensuring no duplicate pairs, can delete
+        counter += 1;
+        if (uniqueShifts.insert(ss).second) {
+          sourceShifts.push_back(ss);
+        }
+      }
+    }
+    LLVM_DEBUG(llvm::dbgs()
+                   << "Total SourceShifts: " << counter << "\n";
+               llvm::dbgs()
+               << "Unique SourceShifts: " << sourceShifts.size() << "\n";);
   }
 
   // Compute the corresponding table of positions after each rotation,
diff --git a/lib/Dialect/TensorExt/Transforms/ShiftScheme.h b/lib/Dialect/TensorExt/Transforms/ShiftScheme.h
