llvm · cachemeifyoucan · Oct 29, 2024
diff --git a/llvm/include/llvm/CAS/ActionCache.h b/llvm/include/llvm/CAS/ActionCache.h
@@ -0,0 +1,94 @@
+//===- llvm/CAS/ActionCache.h -----------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CAS_CASACTIONCACHE_H
+#define LLVM_CAS_CASACTIONCACHE_H
+
+#include "llvm/ADT/StringRef.h"
+#include "llvm/CAS/CASID.h"
+#include "llvm/CAS/CASReference.h"
+#include "llvm/Support/Error.h"
+
+namespace llvm::cas {
+
+class ObjectStore;
+class CASID;
+class ObjectProxy;
+
+/// A key for caching an operation.
+/// It is implemented as a bag of bytes and provides a convenient constructor
+/// for CAS types.
+class CacheKey {
+public:
+  StringRef getKey() const { return Key; }
+
+  // TODO: Support CacheKey other than a CASID but rather any array of bytes.
+  // To do that, ActionCache need to be able to rehash the key into the index,
+  // which then `getOrCompute` method can be used to avoid multiple calls to
+  // has function.
+  CacheKey(const CASID &ID);
+  CacheKey(const ObjectProxy &Proxy);
+  CacheKey(const ObjectStore &CAS, const ObjectRef &Ref);
+
+private:
+  std::string Key;
+};
+
+/// A cache from a key describing an action to the result of doing it.
+///
+/// Actions are expected to be pure (collision is an error).
+class ActionCache {
+  virtual void anchor();
+
+public:
+  /// Get a previously computed result for \p ActionKey.
+  ///
+  /// \param Globally if true it is a hint to the underlying implementation that
+  /// the lookup is profitable to be done on a distributed caching level, not
+  /// just locally. The implementation is free to ignore this flag.
+  Expected<std::optional<CASID>> get(const CacheKey &ActionKey,
+                                     bool Globally = false) const {
+    return getImpl(arrayRefFromStringRef(ActionKey.getKey()), Globally);
+  }
+
+  /// Cache \p Result for the \p ActionKey computation.
+  ///
+  /// \param Globally if true it is a hint to the underlying implementation that
+  /// the association is profitable to be done on a distributed caching level,
+  /// not just locally. The implementation is free to ignore this flag.
+  Error put(const CacheKey &ActionKey, const CASID &Result,
+            bool Globally = false) {
+    assert(Result.getContext().getHashSchemaIdentifier() ==
+               getContext().getHashSchemaIdentifier() &&
+           "Hash schema mismatch");
+    return putImpl(arrayRefFromStringRef(ActionKey.getKey()), Result, Globally);
+  }
+
+  virtual ~ActionCache() = default;
+
+protected:
+  virtual Expected<std::optional<CASID>> getImpl(ArrayRef<uint8_t> ResolvedKey,
+                                                 bool Globally) const = 0;
+
+  virtual Error putImpl(ArrayRef<uint8_t> ResolvedKey, const CASID &Result,
+                        bool Globally) = 0;
+
+  ActionCache(const CASContext &Context) : Context(Context) {}
+
+  const CASContext &getContext() const { return Context; }
+
+private:
+  const CASContext &Context;
+};
+
+/// Create an action cache in memory.
+std::unique_ptr<ActionCache> createInMemoryActionCache();
+
+} // end namespace llvm::cas
+
+#endif // LLVM_CAS_CASACTIONCACHE_H
diff --git a/llvm/lib/CAS/ActionCache.cpp b/llvm/lib/CAS/ActionCache.cpp
@@ -0,0 +1,22 @@
+//===- ActionCache.cpp ------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CAS/ActionCache.h"
+#include "llvm/CAS/CASID.h"
+#include "llvm/CAS/ObjectStore.h"
+
+using namespace llvm;
+using namespace llvm::cas;
+
+void ActionCache::anchor() {}
+
+CacheKey::CacheKey(const CASID &ID) : Key(toStringRef(ID.getHash()).str()) {}
+CacheKey::CacheKey(const ObjectProxy &Proxy)
+    : CacheKey(Proxy.getCAS(), Proxy.getRef()) {}
+CacheKey::CacheKey(const ObjectStore &CAS, const ObjectRef &Ref)
+    : Key(toStringRef(CAS.getID(Ref).getHash())) {}
diff --git a/llvm/lib/CAS/ActionCaches.cpp b/llvm/lib/CAS/ActionCaches.cpp
@@ -0,0 +1,99 @@
+//===- ActionCaches.cpp -----------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "BuiltinCAS.h"
+#include "llvm/ADT/TrieRawHashMap.h"
+#include "llvm/CAS/ActionCache.h"
+#include "llvm/Support/BLAKE3.h"
+
+#define DEBUG_TYPE "action-caches"
+
+using namespace llvm;
+using namespace llvm::cas;
+
+namespace {
+
+using HasherT = BLAKE3;
+using HashType = decltype(HasherT::hash(std::declval<ArrayRef<uint8_t> &>()));
+
+template <size_t Size> class CacheEntry {
+public:
+  CacheEntry() = default;
+  CacheEntry(ArrayRef<uint8_t> Hash) { llvm::copy(Hash, Value.data()); }
+  CacheEntry(const CacheEntry &Entry) { llvm::copy(Entry.Value, Value.data()); }
+  ArrayRef<uint8_t> getValue() const { return Value; }
+
+private:
+  std::array<uint8_t, Size> Value;
+};
+
+class InMemoryActionCache final : public ActionCache {
+public:
+  InMemoryActionCache()
+      : ActionCache(builtin::BuiltinCASContext::getDefaultContext()) {}
+
+  Error putImpl(ArrayRef<uint8_t> ActionKey, const CASID &Result,
+                bool Globally) final;
+  Expected<std::optional<CASID>> getImpl(ArrayRef<uint8_t> ActionKey,
+                                         bool Globally) const final;
+
+private:
+  using DataT = CacheEntry<sizeof(HashType)>;
+  using InMemoryCacheT = ThreadSafeTrieRawHashMap<DataT, sizeof(HashType)>;
+
+  InMemoryCacheT Cache;
+};
+} // end namespace
+
+static std::string hashToString(ArrayRef<uint8_t> Hash) {
+  SmallString<64> Str;
+  toHex(Hash, /*LowerCase=*/true, Str);
+  return Str.str().str();
+}
+
+static Error createResultCachePoisonedError(StringRef Key,
+                                            const CASContext &Context,
+                                            CASID Output,
+                                            ArrayRef<uint8_t> ExistingOutput) {
+  std::string Existing =
+      CASID::create(&Context, toStringRef(ExistingOutput)).toString();
+  return createStringError(std::make_error_code(std::errc::invalid_argument),
+                           "cache poisoned for '" + Key + "' (new='" +
+                               Output.toString() + "' vs. existing '" +
+                               Existing + "')");
+}
+
+Expected<std::optional<CASID>>
+InMemoryActionCache::getImpl(ArrayRef<uint8_t> Key, bool /*Globally*/) const {
+  auto Result = Cache.find(Key);
+  if (!Result)
+    return std::nullopt;
+  return CASID::create(&getContext(), toStringRef(Result->Data.getValue()));
+}
+
+Error InMemoryActionCache::putImpl(ArrayRef<uint8_t> Key, const CASID &Result,
+                                   bool /*Globally*/) {
+  DataT Expected(Result.getHash());
+  const InMemoryCacheT::value_type &Cached = *Cache.insertLazy(
+      Key, [&](auto ValueConstructor) { ValueConstructor.emplace(Expected); });
+
+  const DataT &Observed = Cached.Data;
+  if (Expected.getValue() == Observed.getValue())
+    return Error::success();
+
+  return createResultCachePoisonedError(hashToString(Key), getContext(), Result,
+                                        Observed.getValue());
+}
+
+namespace llvm::cas {
+
+std::unique_ptr<ActionCache> createInMemoryActionCache() {
+  return std::make_unique<InMemoryActionCache>();
+}
+
+} // namespace llvm::cas
diff --git a/llvm/lib/CAS/CMakeLists.txt b/llvm/lib/CAS/CMakeLists.txt
@@ -1,4 +1,6 @@
 add_llvm_component_library(LLVMCAS
+  ActionCache.cpp
+  ActionCaches.cpp
   BuiltinCAS.cpp
   InMemoryCAS.cpp
   ObjectStore.cpp

diff --git a/llvm/unittests/CAS/ActionCacheTest.cpp b/llvm/unittests/CAS/ActionCacheTest.cpp
@@ -0,0 +1,73 @@
+//===- ActionCacheTest.cpp ------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CAS/ActionCache.h"
+#include "CASTestConfig.h"
+#include "llvm/CAS/ObjectStore.h"
+#include "llvm/Testing/Support/Error.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+using namespace llvm::cas;
+
+TEST_P(CASTest, ActionCacheHit) {
+  std::shared_ptr<ObjectStore> CAS = createObjectStore();
+  std::unique_ptr<ActionCache> Cache = createActionCache();
+
+  std::optional<ObjectProxy> ID;
+  ASSERT_THAT_ERROR(CAS->createProxy(std::nullopt, "1").moveInto(ID),
+                    Succeeded());
+  std::optional<CASID> ResultID;
+  ASSERT_THAT_ERROR(Cache->put(*ID, *ID), Succeeded());
+  ASSERT_THAT_ERROR(Cache->get(*ID).moveInto(ResultID), Succeeded());
+  ASSERT_TRUE(ResultID);
+  std::optional<ObjectRef> Result = CAS->getReference(*ResultID);
+  ASSERT_TRUE(Result);
+  ASSERT_EQ(*ID, *Result);
+}
+
+TEST_P(CASTest, ActionCacheMiss) {
+  std::shared_ptr<ObjectStore> CAS = createObjectStore();
+  std::unique_ptr<ActionCache> Cache = createActionCache();
+
+  std::optional<ObjectProxy> ID1, ID2;
+  ASSERT_THAT_ERROR(CAS->createProxy(std::nullopt, "1").moveInto(ID1),
+                    Succeeded());
+  ASSERT_THAT_ERROR(CAS->createProxy(std::nullopt, "2").moveInto(ID2),
+                    Succeeded());
+  ASSERT_THAT_ERROR(Cache->put(*ID1, *ID2), Succeeded());
+  // This is a cache miss for looking up a key doesn't exist.
+  std::optional<CASID> Result1;
+  ASSERT_THAT_ERROR(Cache->get(*ID2).moveInto(Result1), Succeeded());
+  ASSERT_FALSE(Result1);
+
+  ASSERT_THAT_ERROR(Cache->put(*ID2, *ID1), Succeeded());
+  // Cache hit after adding the value.
+  std::optional<CASID> Result2;
+  ASSERT_THAT_ERROR(Cache->get(*ID2).moveInto(Result2), Succeeded());
+  ASSERT_TRUE(Result2);
+  std::optional<ObjectRef> Ref = CAS->getReference(*Result2);
+  ASSERT_TRUE(Ref);
+  ASSERT_EQ(*ID1, *Ref);
+}
+
+TEST_P(CASTest, ActionCacheRewrite) {
+  std::shared_ptr<ObjectStore> CAS = createObjectStore();
+  std::unique_ptr<ActionCache> Cache = createActionCache();
+
+  std::optional<ObjectProxy> ID1, ID2;
+  ASSERT_THAT_ERROR(CAS->createProxy(std::nullopt, "1").moveInto(ID1),
+                    Succeeded());
+  ASSERT_THAT_ERROR(CAS->createProxy(std::nullopt, "2").moveInto(ID2),
+                    Succeeded());
+  ASSERT_THAT_ERROR(Cache->put(*ID1, *ID1), Succeeded());
+  // Writing to the same key with different value is error.
+  ASSERT_THAT_ERROR(Cache->put(*ID1, *ID2), Failed());
+  // Writing the same value multiple times to the same key is fine.
+  ASSERT_THAT_ERROR(Cache->put(*ID1, *ID1), Succeeded());
+}
diff --git a/llvm/unittests/CAS/CASTestConfig.cpp b/llvm/unittests/CAS/CASTestConfig.cpp
@@ -15,7 +15,8 @@ using namespace llvm::cas;
 
 CASTestingEnv createInMemory(int I) {
   std::unique_ptr<ObjectStore> CAS = createInMemoryCAS();
-  return CASTestingEnv{std::move(CAS)};
+  std::unique_ptr<ActionCache> Cache = createInMemoryActionCache();
+  return CASTestingEnv{std::move(CAS), std::move(Cache)};
 }
 
 INSTANTIATE_TEST_SUITE_P(InMemoryCAS, CASTest,

diff --git a/llvm/unittests/CAS/CASTestConfig.h b/llvm/unittests/CAS/CASTestConfig.h
@@ -6,6 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/CAS/ActionCache.h"
 #include "llvm/CAS/ObjectStore.h"
 #include "gtest/gtest.h"
 
@@ -14,6 +15,7 @@
 
 struct CASTestingEnv {
   std::unique_ptr<llvm::cas::ObjectStore> CAS;
+  std::unique_ptr<llvm::cas::ActionCache> Cache;
 };
 
 class CASTest
@@ -25,6 +27,10 @@ class CASTest
     auto TD = GetParam()(++(*NextCASIndex));
     return std::move(TD.CAS);
   }
+  std::unique_ptr<llvm::cas::ActionCache> createActionCache() {
+    auto TD = GetParam()(++(*NextCASIndex));
+    return std::move(TD.Cache);
+  }
   void SetUp() { NextCASIndex = 0; }
   void TearDown() { NextCASIndex = std::nullopt; }
 };

diff --git a/llvm/unittests/CAS/CMakeLists.txt b/llvm/unittests/CAS/CMakeLists.txt
@@ -5,6 +5,7 @@ set(LLVM_LINK_COMPONENTS
   )
 
 add_llvm_unittest(CASTests
+  ActionCacheTest.cpp
   CASTestConfig.cpp
   ObjectStoreTest.cpp
   )