Add a 'Forced' state to lazy_dag_cache and children

This allows one to set a particular key to a known fixed value irrespective of the computations of other nodes. This Forced key acts as a break on updates propagating.

PiperOrigin-RevId: 738582081

Author: allight

xls/passes/BUILD

@@ -1370,6 +1370,7 @@ cc_test(
         "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",

xls/passes/lazy_dag_cache.h

@@ -43,15 +43,20 @@ namespace xls {
 // This class implements a cache with a simple state machine; each key has
 // either no recorded value (kUnknown), a value that may be out-of-date
 // (kUnverified), a value that is valid if all inputs prove to be up-to-date
-// (kInputsUnverified), or a value that is known to be current & correct
-// (kKnown).
+// (kInputsUnverified), a value that is known to be current & correct
+// (kKnown), or information with a forced value (kForced).
 //
 // When a `key` is queried, we query for the values for all of its inputs, and
 // then re-compute the value for `key` if absent or unverified. If `key` was in
 // state kUnverified & this does change its associated value, we mark any direct
 // users that were in state kInputsUnverified as kUnverified, since their inputs
 // have changed.
 //
+// Any node that is in the kForced state will remain in the same state
+// regardless of other changes to the graph. Changing the value a node is Forced
+// to will invalidate all its users and force a recalculation of downstream
+// values.
+//
 // NOTE: If `key` is in state kInputsUnverified after we queried all of its
 //       inputs, then their values did not change, so we have verified that
 //       `key`'s value is up-to-date without having to recompute it! This is the
@@ -93,8 +98,45 @@ class LazyDagCache {
   LazyDagCache<Key, Value>& operator=(LazyDagCache<Key, Value>&& other) =
       delete;
 
+  // Erase all knowledge of the values of all keys.
   void Clear() { cache_.clear(); }
-  void Forget(const Key& key) { cache_.erase(key); }
+  // Erase all knowledge of the value of all keys except for 'Forced' values.
+  void ClearNonForced() {
+    absl::erase_if(cache_, [](const auto& v) {
+      return v.second.state != CacheState::kForced;
+    });
+  }
+
+  // Entirely remove knowledge of this key. This includes erasing any Forced
+  // data.
+  void Forget(const Key& key) {
+    cache_.erase(key);
+    for (const Key& user : provider_->GetUsers(key)) {
+      MarkInputsUnverified(user);
+    }
+  }
+
+  // Set the key as having the immutable, authoritative 'value'.
+  //
+  // *This is a dangerous operation and should be used with care.* It tells the
+  // cache to never call the ComputeValue callback and to instead consider
+  // 'value' to be associated with 'key' now and forever. This knowledge may
+  // only be removed by calling 'Forget' or 'Clear'.
+  void SetForced(const Key& key, std::unique_ptr<Value> value) {
+    cache_.insert_or_assign(key, CacheEntry{.state = CacheState::kForced,
+                                            .value = std::move(value)});
+    MarkUsersUnverified(key);
+  }
+
+  // Set the key as having the immutable, authoritative 'value'.
+  //
+  // *This is a dangerous operation and should be used with care.* It tells the
+  // cache to never call the ComputeValue callback and to instead consider
+  // 'value' to be associated with 'key' now and forever. This knowledge may
+  // only be removed by calling 'Forget' or 'Clear'.
+  void SetForced(const Key& key, Value value) {
+    SetForced(key, std::make_unique<Value>(std::move(value)));
+  }
 
   void AddUnverified(const Key& key, Value value) {
     cache_.insert_or_assign(
@@ -106,6 +148,15 @@ class LazyDagCache {
                                             .value = std::move(value)});
   }
 
+  // Request recomputation of any users of this key.
+  //
+  // Mark as full unverified to force recomputation.
+  void MarkUsersUnverified(const Key& key) {
+    for (const Key& user : provider_->GetUsers(key)) {
+      MarkUnverified(user);
+    }
+  }
+
   void MarkUnverified(const Key& key) {
     auto it = cache_.find(key);
     if (it == cache_.end()) {
@@ -115,6 +166,11 @@ class LazyDagCache {
     if (state == CacheState::kUnverified) {
       return;
     }
+    if (state == CacheState::kForced) {
+      VLOG(1) << "Mark unverified called on forced entry "
+              << provider_->GetName(key);
+      return;
+    }
     state = CacheState::kUnverified;
     for (const Key& user : provider_->GetUsers(key)) {
       if (GetCacheState(user) == CacheState::kKnown) {
@@ -155,7 +211,8 @@ class LazyDagCache {
   Value* QueryValue(const Key& key) {
     // If `key` is already known, return a pointer to the cached value.
     if (auto it = cache_.find(key);
-        it != cache_.end() && it->second.state == CacheState::kKnown) {
+        it != cache_.end() && (it->second.state == CacheState::kKnown ||
+                               it->second.state == CacheState::kForced)) {
       return it->second.value.get();
     }
 
@@ -189,6 +246,7 @@ class LazyDagCache {
       return cached_value.get();
     }
 
+    CHECK_NE(state, CacheState::kForced);
     state = CacheState::kKnown;
     cached_value = std::make_unique<Value>(*std::move(new_value));
 
@@ -208,7 +266,7 @@ class LazyDagCache {
   // all nodes in the DAG.
   absl::Status EagerlyPopulate(absl::Span<const Key> topo_sorted_keys) {
     for (const Key& key : topo_sorted_keys) {
-      if (GetCacheState(key) == CacheState::kKnown) {
+      if (GetNonForcedCacheState(key) == CacheState::kKnown) {
         continue;
       }
       std::vector<const Value*> input_values;
@@ -232,6 +290,8 @@ class LazyDagCache {
   // consistent regardless. This is an expensive operation, intended for use in
   // tests. `topo_sorted_keys` must be a topological sort of the keys for all
   // nodes in the DAG.
+  //
+  // Note that Forced values are always considered consistent.
   absl::Status CheckConsistency(absl::Span<const Key> topo_sorted_keys) const;
 
   enum class CacheState : uint8_t {
@@ -241,6 +301,14 @@ class LazyDagCache {
     kUnverified,
     kInputsUnverified,
     kKnown,
+    // A value has been provided via external information and should be
+    // considered authoritatively known.
+    //
+    // This node can never be put to unverified state.
+    //
+    // It is possible that the value this key is forced to is one that cannot be
+    // arrived at through the normal update sequence.
+    kForced,
   };
 
   template <typename Sink>
@@ -258,6 +326,9 @@ class LazyDagCache {
       case CacheState::kKnown:
         absl::Format(&sink, "KNOWN");
         return;
+      case CacheState::kForced:
+        absl::Format(&sink, "FORCED");
+        return;
     }
     LOG(FATAL) << "Unknown CacheState: " << static_cast<int>(state);
   }
@@ -272,6 +343,14 @@ class LazyDagCache {
     }
     return it->second.state;
   }
+  // Get the cache state with forced values being considered kKnown
+  CacheState GetNonForcedCacheState(const Key& key) const {
+    CacheState s = GetCacheState(key);
+    if (s == CacheState::kForced) {
+      return CacheState::kKnown;
+    }
+    return s;
+  }
   const Value* GetCachedValue(const Key& key) const {
     auto it = cache_.find(key);
     if (it == cache_.end() || it->second.value == nullptr) {
@@ -320,9 +399,10 @@ absl::Status LazyDagCache<Key, Value>::CheckConsistency(
 
     if (state == CacheState::kKnown) {
       for (const Key& input : provider_->GetInputs(key)) {
-        XLS_RET_CHECK_EQ(GetCacheState(input), CacheState::kKnown)
+        XLS_RET_CHECK_EQ(GetNonForcedCacheState(input), CacheState::kKnown)
             << "Non-KNOWN input for KNOWN key " << provider_->GetName(key)
-            << ": " << provider_->GetName(input);
+            << ": " << provider_->GetName(input) << " (input is "
+            << GetCacheState(input) << ")";
       }
     }
     if (state == CacheState::kInputsUnverified) {
@@ -333,6 +413,8 @@ absl::Status LazyDagCache<Key, Value>::CheckConsistency(
       }
     }
 
+    // NB state FORCED & UNVERIFIED has no requirements on its inputs.
+
     if (absl::c_any_of(provider_->GetInputs(key), [&](const Key& input) {
           return !correct_values.contains(input);
         })) {
@@ -341,6 +423,12 @@ absl::Status LazyDagCache<Key, Value>::CheckConsistency(
       continue;
     }
 
+    if (state == CacheState::kForced) {
+      // Forced values are definitionally correct.
+      correct_values.insert(key);
+      continue;
+    }
+
     std::vector<const Value*> input_values;
     absl::Span<const Key> inputs = provider_->GetInputs(key);
     input_values.reserve(inputs.size());

xls/passes/lazy_node_info.h

@@ -55,8 +55,8 @@ namespace xls {
 // This class implements a cache with a simple state machine; each node has
 // either no recorded information (kUnknown), information that may be
 // out-of-date (kUnverified), information that is valid if all inputs prove to
-// be up-to-date (kInputsUnverified), or information that is known to be current
-// & correct (kKnown).
+// be up-to-date (kInputsUnverified), information that is known to be current
+// & correct (kKnown), or information with a forced value (kForced).
 //
 // On populating with a FunctionBase, this class listens for change events. When
 // a node changes, we mark its information as potentially out-of-date
@@ -71,6 +71,12 @@ namespace xls {
 // information, we mark any direct users that were in state kInputsUnverified as
 // kUnverified, since their inputs have changed.
 //
+// Any node that is in the kForced state will remain in the same state
+// regardless of other changes to the graph. Changing the value a node is Forced
+// to will invalidate all its users.
+//
+// A node may not have both 'given' data and 'forced' data at the same time.
+//
 // NOTE: If `n` is in state kInputsUnverified after we queried all of its
 //       operands, then their values did not change, so we have verified that
 //       `n`'s information is up-to-date without having to recompute it! This is
@@ -177,13 +183,51 @@ class LazyNodeInfo
     return rf;
   }
 
+  // Set the node to a single immutable forced value.
+  //
+  // This is different from Givens since it is not combined with the calculated
+  // values from earlier in the tree but instead considered a-priori known.
+  //
+  // Note that any forced value may not have a given associated with it as the
+  // given value will be ignored.
+  //
+  // Care should be taken when using this since existing information is utterly
+  // ignored. In general AddGiven is a better choice.
+  absl::StatusOr<ReachedFixpoint> SetForced(Node* node,
+                                            LeafTypeTree<Info> forced_ltt) {
+    XLS_RET_CHECK(givens_.find(node) == givens_.end())
+        << node << " already has given information.";
+    if (cache_.GetCacheState(node) == CacheState::kForced &&
+        *cache_.GetCachedValue(node) == forced_ltt) {
+      return ReachedFixpoint::Unchanged;
+    }
+    cache_.SetForced(node, std::move(forced_ltt));
+    return ReachedFixpoint::Changed;
+  }
+
+  absl::StatusOr<ReachedFixpoint> RemoveForced(Node* node) {
+    XLS_RET_CHECK_EQ(cache_.GetCacheState(node), CacheState::kForced)
+        << node << " has no forced value associated with it.";
+    if (cache_.GetCacheState(node) != CacheState::kForced) {
+      return ReachedFixpoint::Unchanged;
+    }
+    cache_.Forget(node);
+    return ReachedFixpoint::Changed;
+  }
+
   // Bind the node info to the given function.
   absl::StatusOr<ReachedFixpoint> Attach(FunctionBase* f) {
     return AttachWithGivens(f, {});
   }
 
+  // Set the value 'given' as being assumed for the given node. This data is
+  // combined with the already calculated data. If one needs to set the state
+  // directly to the given value use SetForced instead. A node may not have both
+  // 'Forced' and 'Given' data associated with it at the same time.
   absl::StatusOr<ReachedFixpoint> AddGiven(Node* node,
                                            LeafTypeTree<Info> given_ltt) {
+    XLS_RET_CHECK_NE(cache_.GetCacheState(node), CacheState::kForced)
+        << node << " has a forced value associated with it.";
     auto it = givens_.find(node);
     if (it == givens_.end()) {
       givens_.emplace(node, std::move(given_ltt));
@@ -207,9 +251,13 @@ class LazyNodeInfo
     cache_.MarkUnverified(node);
     return ReachedFixpoint::Changed;
   }
-  ReachedFixpoint ReplaceGiven(Node* node, LeafTypeTree<Info> given) {
+
+  absl::StatusOr<ReachedFixpoint> ReplaceGiven(Node* node,
+                                               LeafTypeTree<Info> given) {
     auto it = givens_.find(node);
     if (it == givens_.end()) {
+      XLS_RET_CHECK_NE(cache_.GetCacheState(node), CacheState::kForced)
+          << node << " has a forced value associated with it.";
       givens_.emplace(node, std::move(given));
       cache_.MarkUnverified(node);
       return ReachedFixpoint::Changed;
@@ -287,6 +335,8 @@ class LazyNodeInfo
   // lazy query engines, checks that the current state of the cache is correct
   // where expected & consistent regardless. This is an expensive operation,
   // intended for use in tests.
+  //
+  // Note that Forced values are always considered consistent.
   absl::Status CheckCacheConsistency() const {
     XLS_RET_CHECK(f_ != nullptr) << "Unattached info";
     return cache_.CheckConsistency(TopoSort(f_));

xls/passes/lazy_query_engine.h

@@ -136,6 +136,8 @@ class LazyQueryEngine : public QueryEngine {
     return node->function_base() == info_.bound_function();
   }
 
+  // Check that the query engine is consistent. Note that Forced values are
+  // always considered consistent.
   absl::Status CheckConsistency() const override {
     return info_.CheckCacheConsistency();
   }
@@ -172,6 +174,25 @@ class LazyQueryEngine : public QueryEngine {
   }
   ReachedFixpoint RemoveGiven(Node* node) { return info_.RemoveGiven(node); }
 
+  // Set the node to a single immutable forced value.
+  //
+  // This is different from Givens since it is not combined with the calculated
+  // values from earlier in the tree but instead considered a-priori known.
+  //
+  // Note that any forced value may not have a given associated with it as the
+  // given value will be ignored.
+  //
+  // Care should be taken when using this since existing information is utterly
+  // ignored. In general AddGiven is a better choice.
+  absl::StatusOr<ReachedFixpoint> SetForced(Node* node,
+                                            LeafTypeTree<Info> forced_ltt) {
+    return info_.SetForced(node, std::move(forced_ltt));
+  }
+  // Removed forced information.
+  absl::StatusOr<ReachedFixpoint> RemoveForced(Node* node) {
+    return info_.RemoveForced(node);
+  }
+
  protected:
   virtual LeafTypeTree<Info> ComputeInfo(
       Node* node,