Metal: remove cached pipeline states with deleted programs (#8243)

google · Nov 7, 2024 · b976839 · b976839
1 parent 3e02d50
commit b976839
Show file tree

Hide file tree

Showing 5 changed files with 39 additions and 3 deletions.
diff --git a/filament/backend/src/metal/MetalDriver.mm b/filament/backend/src/metal/MetalDriver.mm
@@ -797,9 +797,20 @@
 }
 
 void MetalDriver::destroyProgram(Handle<HwProgram> ph) {
-    if (ph) {
-        destruct_handle<MetalProgram>(ph);
+    if (UTILS_UNLIKELY(!ph)) {
+        return;
+    }
+    // Remove any cached pipeline states that refer to these programs.
+    auto* metalProgram = handle_cast<MetalProgram>(ph);
+    const auto& functions = metalProgram->getFunctionsIfPresent();
+    if (UTILS_LIKELY(functions.isRaster())) { // only raster pipelines are cached
+        const auto& raster = functions.getRasterFunctions();
+        mContext->pipelineStateCache.removeIf([&](const MetalPipelineState& state) {
+            const auto& [fragment, vertex] = raster;
+            return state.fragmentFunction == fragment || state.vertexFunction == vertex;
+        });
     }
+    destruct_handle<MetalProgram>(ph);
 }
 
 void MetalDriver::destroyTexture(Handle<HwTexture> th) {

diff --git a/filament/backend/src/metal/MetalHandles.h b/filament/backend/src/metal/MetalHandles.h
@@ -202,6 +202,7 @@ class MetalProgram : public HwProgram {
     MetalProgram(MetalContext& context, Program&& program) noexcept;
 
     const MetalShaderCompiler::MetalFunctionBundle& getFunctions();
+    const MetalShaderCompiler::MetalFunctionBundle& getFunctionsIfPresent() const;
 
 private:
     void initialize();

diff --git a/filament/backend/src/metal/MetalHandles.mm b/filament/backend/src/metal/MetalHandles.mm
@@ -495,6 +495,10 @@ static void func(void* user) {
     return mFunctionBundle;
 }
 
+const MetalShaderCompiler::MetalFunctionBundle& MetalProgram::getFunctionsIfPresent() const {
+    return mFunctionBundle;
+}
+
 void MetalProgram::initialize() {
     if (!mToken) {
         return;

diff --git a/filament/backend/src/metal/MetalShaderCompiler.h b/filament/backend/src/metal/MetalShaderCompiler.h
@@ -83,6 +83,10 @@ class MetalShaderCompiler {
             return std::get<Compute>(mPrograms);
         }
 
+        bool isRaster() const { return std::holds_alternative<Raster>(mPrograms); }
+
+        bool isCompute() const { return std::holds_alternative<Compute>(mPrograms); }
+
         static MetalFunctionBundle none() {
             return MetalFunctionBundle(None{});
         }

diff --git a/filament/backend/src/metal/MetalState.h b/filament/backend/src/metal/MetalState.h
@@ -28,7 +28,9 @@
 
 #include <memory>
 #include <tsl/robin_map.h>
+
 #include <utils/Hash.h>
+#include <utils/Invocable.h>
 
 namespace filament {
 namespace backend {
@@ -160,6 +162,8 @@ template<typename StateType,
          typename HashFn = utils::hash::MurmurHashFn<StateType>>
 class StateCache {
 
+    using MapType = tsl::robin_map<StateType, MetalType, HashFn>;
+
 public:
 
     StateCache() = default;
@@ -169,6 +173,18 @@ class StateCache {
 
     void setDevice(id<MTLDevice> device) noexcept { mDevice = device; }
 
+    void removeIf(utils::Invocable<bool(const StateType&)> fn) noexcept {
+        typename MapType::const_iterator it = mStateCache.begin();
+        while (it != mStateCache.end()) {
+            const auto& [key, _] = *it;
+            if (UTILS_UNLIKELY(fn(key))) {
+                it = mStateCache.erase(it);
+            } else {
+                ++it;
+            }
+        }
+    }
+
     MetalType getOrCreateState(const StateType& state) noexcept {
         assert_invariant(mDevice);
 
@@ -196,7 +212,7 @@ class StateCache {
     StateCreator creator;
     id<MTLDevice> mDevice = nil;
 
-    tsl::robin_map<StateType, MetalType, HashFn> mStateCache;
+    MapType mStateCache;
 
 };